[SSH] Support IDE path override for SSH

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

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

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,27 +81,56 @@ 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
 	}
 	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 %s 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
 

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-<uuid> on every serverless instance, so a "<user>@" prefix would give each
 // connect a different authority and VS Code would remember a separate

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.