From 9b23a1f91c14440c1b7f2c71c172c8b33b1f47b9 Mon Sep 17 00:00:00 2001 From: Russell Clarey Date: Tue, 3 Nov 2026 15:49:11 +0100 Subject: [PATCH 1/2] support IDE path override for SSH --- experimental/ssh/internal/client/client.go | 13 +++- .../internal/client/client_internal_test.go | 7 ++ experimental/ssh/internal/vscode/run.go | 65 ++++++++++++---- experimental/ssh/internal/vscode/run_test.go | 77 ++++++++++++++++++- libs/telemetry/protos/ssh_tunnel.go | 5 ++ 5 files changed, 147 insertions(+), 20 deletions(-) diff --git a/experimental/ssh/internal/client/client.go b/experimental/ssh/internal/client/client.go index 80142d54d4c..c35e6dde34c 100644 --- a/experimental/ssh/internal/client/client.go +++ b/experimental/ssh/internal/client/client.go @@ -365,8 +365,8 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt } if opts.IDE != "" && !opts.ProxyMode { - if err := vscode.CheckIDECommand(opts.IDE); err != nil { - outcome.errorCategory = protos.SshTunnelErrorCategoryIDECommandNotOnPath + if err := vscode.CheckIDECommand(ctx, opts.IDE); err != nil { + outcome.errorCategory = ideCommandErrorCategory(err) return err } if err := vscode.CheckIDESSHExtension(ctx, opts.IDE, opts.AutoApprove); err != nil { @@ -1493,6 +1493,15 @@ type connectOutcome struct { ctxErr error } +// ideCommandErrorCategory tells a broken IDE command override apart from a command missing on +// PATH: the first points at whatever set the override, the second at the user's machine. +func ideCommandErrorCategory(err error) protos.SshTunnelErrorCategory { + if errors.Is(err, vscode.ErrIDECommandOverrideInvalid) { + return protos.SshTunnelErrorCategoryIDECommandOverrideInvalid + } + return protos.SshTunnelErrorCategoryIDECommandNotOnPath +} + // sshExtensionErrorCategory attributes a Remote SSH extension check failure to the outcome that // caused it. The four are kept apart because they imply different fixes, and because only the // first two can occur under --auto-approve, which the IDE button always passes -- so a shift diff --git a/experimental/ssh/internal/client/client_internal_test.go b/experimental/ssh/internal/client/client_internal_test.go index 1db7adbd88a..82b6f14d660 100644 --- a/experimental/ssh/internal/client/client_internal_test.go +++ b/experimental/ssh/internal/client/client_internal_test.go @@ -695,6 +695,13 @@ func TestSshExtensionErrorCategory(t *testing.T) { } } +func TestIDECommandErrorCategory(t *testing.T) { + assert.Equal(t, protos.SshTunnelErrorCategoryIDECommandOverrideInvalid, + ideCommandErrorCategory(fmt.Errorf("wrapped: %w", vscode.ErrIDECommandOverrideInvalid))) + assert.Equal(t, protos.SshTunnelErrorCategoryIDECommandNotOnPath, + ideCommandErrorCategory(errors.New("command not found on PATH"))) +} + func TestProxySessionEndCategory(t *testing.T) { tests := []struct { name string diff --git a/experimental/ssh/internal/vscode/run.go b/experimental/ssh/internal/vscode/run.go index 912e35a5f17..db51857424f 100644 --- a/experimental/ssh/internal/vscode/run.go +++ b/experimental/ssh/internal/vscode/run.go @@ -6,9 +6,11 @@ import ( "fmt" "os" "os/exec" + "path/filepath" "strings" "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/log" "golang.org/x/mod/semver" ) @@ -79,6 +81,16 @@ var cursorIDE = ideDescriptor{ LaunchArgs: []string{"--classic"}, } +// IDECommandEnv overrides the IDE command with an absolute path. The VS Code extension sets it +// to the running editor's own CLI, so connecting doesn't depend on the command being on PATH. It +// replaces the command of whichever IDE --ide selects, so it must point at that IDE's CLI. +const IDECommandEnv = "DATABRICKS_SSH_IDE_COMMAND" + +// ErrIDECommandOverrideInvalid is returned when the IDECommandEnv override can't be run. It is +// kept apart from a command missing on PATH because it points at whatever set the variable +// rather than at the user's PATH. +var ErrIDECommandOverrideInvalid = errors.New("IDE command override is not runnable") + func getIDE(option string) ideDescriptor { if option == CursorOption { return cursorIDE @@ -86,20 +98,39 @@ func getIDE(option string) ideDescriptor { return vsCodeIDE } -// CheckIDECommand verifies the IDE CLI command is available on PATH. -func CheckIDECommand(option string) error { +// ideCommand returns the command to run the IDE CLI with: the IDECommandEnv override if set, +// otherwise the IDE's command name to look up on PATH. +func ideCommand(ctx context.Context, ide ideDescriptor) string { + if command := env.Get(ctx, IDECommandEnv); command != "" { + return command + } + return ide.Command +} + +// CheckIDECommand verifies the IDE CLI command is available. +func CheckIDECommand(ctx context.Context, option string) error { ide := getIDE(option) - if _, err := exec.LookPath(ide.Command); err != nil { - return fmt.Errorf( - "%q command not found on PATH. To fix this:\n"+ - "1. Install %s from %s\n"+ - "2. Open the Command Palette (Cmd+Shift+P / Ctrl+Shift+P) and run \"Shell Command: Install '%s' command\"\n"+ - "3. Restart your terminal", - ide.Command, ide.Name, ide.InstallURL, ide.Command, - ) + if command := env.Get(ctx, IDECommandEnv); command != "" { + if !filepath.IsAbs(command) { + return fmt.Errorf("%w: %s must be an absolute path, got %q", ErrIDECommandOverrideInvalid, IDECommandEnv, command) + } + if _, err := exec.LookPath(command); err != nil { + return fmt.Errorf("%w: %s command %q from %s: %w", ErrIDECommandOverrideInvalid, ide.Name, command, IDECommandEnv, err) + } + return nil } - return nil + + if _, err := exec.LookPath(ide.Command); err == nil { + return nil + } + return fmt.Errorf( + "%q command not found on PATH. To fix this:\n"+ + "1. Install %s from %s\n"+ + "2. Open the Command Palette (Cmd+Shift+P / Ctrl+Shift+P) and run \"Shell Command: Install '%s' command\"\n"+ + "3. Restart your terminal", + ide.Command, ide.Name, ide.InstallURL, ide.Command, + ) } // parseExtensionVersion finds the version of the given extension in the output @@ -157,7 +188,9 @@ func consentErrorf(sentinel error, format string, args ...any) error { func CheckIDESSHExtension(ctx context.Context, option string, autoApprove bool) error { ide := getIDE(option) - out, err := exec.CommandContext(ctx, ide.Command, "--list-extensions", "--show-versions").Output() + command := ideCommand(ctx, ide) + + out, err := exec.CommandContext(ctx, command, "--list-extensions", "--show-versions").Output() if err != nil { return fmt.Errorf("%w in %s: %w", ErrSSHExtensionListFailed, ide.Name, err) } @@ -179,7 +212,7 @@ func CheckIDESSHExtension(ctx context.Context, option string, autoApprove bool) if !cmdio.IsPromptSupported(ctx) { return consentErrorf(ErrSSHExtensionInstallUnavailable, "%s Install it with: %s --install-extension %s, or pass --auto-approve", - msg, ide.Command, ide.SSHExtensionID) + msg, command, ide.SSHExtensionID) } shouldInstall, err := cmdio.AskYesOrNo(ctx, msg+" Would you like to install it?") @@ -189,14 +222,14 @@ func CheckIDESSHExtension(ctx context.Context, option string, autoApprove bool) if !shouldInstall { return consentErrorf(ErrSSHExtensionInstallDeclined, "%s Install it with: %s --install-extension %s", - msg, ide.Command, ide.SSHExtensionID) + msg, command, ide.SSHExtensionID) } } else { cmdio.LogString(ctx, msg+" Installing automatically (--auto-approve).") } cmdio.LogString(ctx, fmt.Sprintf("Installing %q...", ide.SSHExtensionName)) - installCmd := exec.CommandContext(ctx, ide.Command, "--install-extension", ide.SSHExtensionID, "--force") + installCmd := exec.CommandContext(ctx, command, "--install-extension", ide.SSHExtensionID, "--force") installCmd.Stdout = os.Stdout installCmd.Stderr = os.Stderr if err := installCmd.Run(); err != nil { @@ -229,7 +262,7 @@ func LaunchIDE(ctx context.Context, ideOption, connectionName, databricksUserNam log.Infof(ctx, "Launching %s with args: %v", ideOption, args) - ideCmd := exec.CommandContext(ctx, ide.Command, args...) + ideCmd := exec.CommandContext(ctx, ideCommand(ctx, ide), args...) ideCmd.Stdout = os.Stdout ideCmd.Stderr = os.Stderr diff --git a/experimental/ssh/internal/vscode/run_test.go b/experimental/ssh/internal/vscode/run_test.go index 839ef931a3b..d6a60d28224 100644 --- a/experimental/ssh/internal/vscode/run_test.go +++ b/experimental/ssh/internal/vscode/run_test.go @@ -46,7 +46,7 @@ func TestCheckIDECommand_Missing(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - err := CheckIDECommand(tt.ide) + err := CheckIDECommand(t.Context(), tt.ide) require.Error(t, err) assert.Contains(t, err.Error(), tt.wantErrMsg) }) @@ -86,12 +86,44 @@ func TestCheckIDECommand_Found(t *testing.T) { err := os.WriteFile(fakePath, []byte("#!/bin/sh\n"), 0o755) require.NoError(t, err) - err = CheckIDECommand(tt.ide) + err = CheckIDECommand(t.Context(), tt.ide) assert.NoError(t, err) }) } } +func TestCheckIDECommand_Override(t *testing.T) { + // The override must be used even though nothing is on PATH. + t.Setenv("PATH", t.TempDir()) + binDir := t.TempDir() + createFakeIDEExecutable(t, binDir, "code", "") + t.Setenv(IDECommandEnv, fakeIDECommandPath(binDir, "code")) + + assert.NoError(t, CheckIDECommand(t.Context(), VSCodeOption)) +} + +func TestCheckIDECommand_OverrideMissing(t *testing.T) { + missing := filepath.Join(t.TempDir(), "code") + t.Setenv(IDECommandEnv, missing) + + err := CheckIDECommand(t.Context(), VSCodeOption) + require.ErrorIs(t, err, ErrIDECommandOverrideInvalid) + assert.Contains(t, err.Error(), IDECommandEnv) + assert.Contains(t, err.Error(), missing) +} + +func TestCheckIDECommand_OverrideNotAbsolute(t *testing.T) { + // A bare name is rejected even when it resolves on PATH. + binDir := t.TempDir() + t.Setenv("PATH", binDir) + createFakeIDEExecutable(t, binDir, "code", "") + t.Setenv(IDECommandEnv, "code") + + err := CheckIDECommand(t.Context(), VSCodeOption) + require.ErrorIs(t, err, ErrIDECommandOverrideInvalid) + assert.Contains(t, err.Error(), "must be an absolute path") +} + func TestParseExtensionVersion(t *testing.T) { tests := []struct { name string @@ -233,6 +265,14 @@ func createFakeIDEExecutable(t *testing.T, dir, command, output string) { } } +// fakeIDECommandPath returns the path of a command written by the fake IDE helpers. +func fakeIDECommandPath(dir, command string) string { + if runtime.GOOS == "windows" { + command += ".cmd" + } + return filepath.Join(dir, command) +} + // createFailingIDEExecutable writes a fake IDE command that exits non-zero for every // invocation, so "--list-extensions" fails and the check never learns what is installed. func createFailingIDEExecutable(t *testing.T, dir, command string) { @@ -343,6 +383,39 @@ func TestCheckIDESSHExtension_AutoApproveMissing_Installs(t *testing.T) { assert.NoError(t, err) } +func TestCheckIDESSHExtension_UsesOverride(t *testing.T) { + t.Setenv("PATH", t.TempDir()) + binDir := t.TempDir() + createFakeIDEExecutable(t, binDir, "code", "ms-vscode-remote.remote-ssh@0.123.0\n") + t.Setenv(IDECommandEnv, fakeIDECommandPath(binDir, "code")) + ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) + + assert.NoError(t, CheckIDESSHExtension(ctx, VSCodeOption, false)) +} + +func TestCheckIDESSHExtension_OverrideInInstallHint(t *testing.T) { + t.Setenv("PATH", t.TempDir()) + binDir := t.TempDir() + createFakeIDEExecutable(t, binDir, "code", "ms-python.python@2024.1.1\n") + command := fakeIDECommandPath(binDir, "code") + t.Setenv(IDECommandEnv, command) + ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) + + err := CheckIDESSHExtension(ctx, VSCodeOption, false) + require.ErrorIs(t, err, ErrSSHExtensionInstallUnavailable) + assert.Contains(t, err.Error(), command+" --install-extension") +} + +func TestLaunchIDE_UsesOverride(t *testing.T) { + // Nothing is on PATH, so the launch only succeeds if it runs the override. + t.Setenv("PATH", t.TempDir()) + binDir := t.TempDir() + createFakeIDEExecutable(t, binDir, "code", "") + t.Setenv(IDECommandEnv, fakeIDECommandPath(binDir, "code")) + + assert.NoError(t, LaunchIDE(t.Context(), VSCodeOption, "databricks-cpu-7f189c39", "me@example.com")) +} + // The authority must be the host alias alone. The remote OS user is a fresh // spark- on every serverless instance, so a "@" prefix would give each // connect a different authority and VS Code would remember a separate diff --git a/libs/telemetry/protos/ssh_tunnel.go b/libs/telemetry/protos/ssh_tunnel.go index 708a453ac63..6b54b6e5960 100644 --- a/libs/telemetry/protos/ssh_tunnel.go +++ b/libs/telemetry/protos/ssh_tunnel.go @@ -39,6 +39,11 @@ const ( // condition rather than a transient failure, so it is distinguished from the rest. SshTunnelErrorCategoryIDECommandNotOnPath SshTunnelErrorCategory = "IDE_COMMAND_NOT_ON_PATH" + // The IDE command override (DATABRICKS_SSH_IDE_COMMAND, set by the VS Code extension) names + // a command that can't be run. Distinct from IDE_COMMAND_NOT_ON_PATH because it points at + // whatever set the override rather than at the user's PATH. + SshTunnelErrorCategoryIDECommandOverrideInvalid SshTunnelErrorCategory = "IDE_COMMAND_OVERRIDE_INVALID" + // The IDE's installed-extension list could not be read, so whether the Remote SSH // extension was present is unknown. Distinct from the install failures below because it // says nothing about the extension itself, only that the check could not run. From a286464d9bc27b2941070ea06f3d1de462cc379a Mon Sep 17 00:00:00 2001 From: Russell Clarey Date: Tue, 3 Nov 2026 16:19:10 +0100 Subject: [PATCH 2/2] fix error quoting for windows --- experimental/ssh/internal/vscode/run.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/experimental/ssh/internal/vscode/run.go b/experimental/ssh/internal/vscode/run.go index db51857424f..bb006a7c626 100644 --- a/experimental/ssh/internal/vscode/run.go +++ b/experimental/ssh/internal/vscode/run.go @@ -116,7 +116,7 @@ func CheckIDECommand(ctx context.Context, option string) error { return fmt.Errorf("%w: %s must be an absolute path, got %q", ErrIDECommandOverrideInvalid, IDECommandEnv, command) } if _, err := exec.LookPath(command); err != nil { - return fmt.Errorf("%w: %s command %q from %s: %w", ErrIDECommandOverrideInvalid, ide.Name, command, IDECommandEnv, err) + return fmt.Errorf("%w: %s command %s from %s: %w", ErrIDECommandOverrideInvalid, ide.Name, command, IDECommandEnv, err) } return nil }