Skip prompts in ssh connect proxy mode and ssh server command (#3634)

## Changes
- Skip prompts in ssh connect proxy mode and ssh server command
- Add profile option to the ProxyCommand when spawning ssh client from
`databricks ssh connect` (and the setup logic already adds it to the
config).
- Skip loading bundles in ssh server command, so the bundle doesn't
override the auth that's setup by the job logic that executes this
command (or worse, the server can fail to start if the bundle happens to
be incorrect)


## Tests
Only manual

Author: ilia-db

experimental/ssh/cmd/connect.go

@@ -46,7 +46,17 @@ the SSH server and handling the connection proxy.
 	cmd.Flags().StringVar(&releasesDir, "releases-dir", "", "Directory for local SSH tunnel development releases")
 	cmd.Flags().MarkHidden("releases-dir")
 
-	cmd.PreRunE = root.MustWorkspaceClient
+	cmd.PreRunE = func(cmd *cobra.Command, args []string) error {
+		// CLI in the proxy mode is executed by the ssh client and can't prompt for input
+		if proxyMode {
+			cmd.SetContext(root.SkipPrompt(cmd.Context()))
+		}
+		// We want to avoid the situation where the connect command works because it pulls the auth config from a bundle,
+		// but fails if it's executed outside of it (which will happen when using remote development IDE features).
+		cmd.SetContext(root.SkipLoadBundle(cmd.Context()))
+		return root.MustWorkspaceClient(cmd, args)
+	}
+
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
 		ctx := cmd.Context()
 		wsClient := cmdctx.WorkspaceClient(ctx)
@@ -62,6 +72,7 @@ the SSH server and handling the connection proxy.
 			ClientPublicKeyName: defaultClientPublicKeyName,
 			ServerTimeout:       serverTimeout,
 			AutoStartCluster:    autoStartCluster,
+			Profile:             wsClient.Config.Profile,
 		}
 		return client.Run(ctx, wsClient, opts)
 	}

experimental/ssh/cmd/server.go

@@ -41,7 +41,16 @@ and proxies them to local SSH daemon processes.
 	cmd.Flags().DurationVar(&shutdownDelay, "shutdown-delay", defaultShutdownDelay, "Delay before shutting down after no pings from clients")
 	cmd.Flags().StringVar(&version, "version", "", "Client version of the Databricks CLI")
 
-	cmd.PreRunE = root.MustWorkspaceClient
+	cmd.PreRunE = func(cmd *cobra.Command, args []string) error {
+		// The server can be executed under a directory with an invalid bundle configuration.
+		// We do not want to error out in this case.
+		// The auth is setup by the job logic that executes this command.
+		cmd.SetContext(root.SkipLoadBundle(cmd.Context()))
+		// The command should be executed in a non-interactive environment, but let's be explicit about no prompts.
+		cmd.SetContext(root.SkipPrompt(cmd.Context()))
+		return root.MustWorkspaceClient(cmd, args)
+	}
+
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
 		ctx := cmd.Context()
 		wsc := cmdctx.WorkspaceClient(ctx)

experimental/ssh/cmd/setup.go

@@ -34,7 +34,13 @@ an SSH host configuration to your SSH config file.
 	cmd.Flags().StringVar(&sshConfigPath, "ssh-config", "", "Path to SSH config file (default ~/.ssh/config)")
 	cmd.Flags().DurationVar(&shutdownDelay, "shutdown-delay", defaultShutdownDelay, "SSH server will terminate after this delay if there are no active connections")
 
-	cmd.PreRunE = root.MustWorkspaceClient
+	cmd.PreRunE = func(cmd *cobra.Command, args []string) error {
+		// We want to avoid the situation where the setup command works because it pulls the auth config from a bundle,
+		// but later on the `ssh host-name` command fails when executed outside of the bundle directory.
+		cmd.SetContext(root.SkipLoadBundle(cmd.Context()))
+		return root.MustWorkspaceClient(cmd, args)
+	}
+
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
 		ctx := cmd.Context()
 		client := cmdctx.WorkspaceClient(ctx)

experimental/ssh/internal/client/client.go

@@ -19,6 +19,7 @@ import (
 
 	"github.com/databricks/cli/experimental/ssh/internal/keys"
 	"github.com/databricks/cli/experimental/ssh/internal/proxy"
+	"github.com/databricks/cli/experimental/ssh/internal/setup"
 	sshWorkspace "github.com/databricks/cli/experimental/ssh/internal/workspace"
 	"github.com/databricks/cli/internal/build"
 	"github.com/databricks/cli/libs/cmdio"
@@ -61,6 +62,8 @@ type ClientOptions struct {
 	ClientPublicKeyName string
 	// If true, the CLI will attempt to start the cluster if it is not running.
 	AutoStartCluster bool
+	// Optional auth profile name. If present, will be added as --profile flag to the ProxyCommand while spawning ssh client.
+	Profile string
 	// Additional arguments to pass to the SSH client in the non proxy mode.
 	AdditionalArgs []string
 }
@@ -231,14 +234,11 @@ func submitSSHTunnelJob(ctx context.Context, client *databricks.WorkspaceClient,
 }
 
 func spawnSSHClient(ctx context.Context, userName, privateKeyPath string, serverPort int, opts ClientOptions) error {
-	executablePath, err := os.Executable()
+	proxyCommand, err := setup.GenerateProxyCommand(opts.ClusterID, opts.AutoStartCluster, opts.ShutdownDelay, opts.Profile, userName, serverPort, opts.HandoverTimeout)
 	if err != nil {
-		return fmt.Errorf("failed to get current executable path: %w", err)
+		return fmt.Errorf("failed to generate ProxyCommand: %w", err)
 	}
 
-	proxyCommand := fmt.Sprintf("%s ssh connect --proxy --cluster=%s --handover-timeout=%s --metadata=%s,%d --auto-start-cluster=%t",
-		executablePath, opts.ClusterID, opts.HandoverTimeout.String(), userName, serverPort, opts.AutoStartCluster)
-
 	sshArgs := []string{
 		"-l", userName,
 		"-i", privateKeyPath,

experimental/ssh/internal/setup/setup.go

@@ -7,6 +7,7 @@ import (
 	"os"
 	"path/filepath"
 	"regexp"
+	"strconv"
 	"strings"
 	"time"
 
@@ -55,20 +56,39 @@ func resolveConfigPath(configPath string) (string, error) {
 	return filepath.Join(homeDir, ".ssh", "config"), nil
 }
 
+func GenerateProxyCommand(clusterId string, autoStartCluster bool, shutdownDelay time.Duration, profile, userName string, serverPort int, handoverTimeout time.Duration) (string, error) {
+	executablePath, err := os.Executable()
+	if err != nil {
+		return "", fmt.Errorf("failed to get current executable path: %w", err)
+	}
+
+	proxyCommand := fmt.Sprintf("%q ssh connect --proxy --cluster=%s --auto-start-cluster=%t --shutdown-delay=%s",
+		executablePath, clusterId, autoStartCluster, shutdownDelay.String())
+
+	if userName != "" && serverPort != 0 {
+		proxyCommand += " --metadata=" + userName + "," + strconv.Itoa(serverPort)
+	}
+
+	if handoverTimeout > 0 {
+		proxyCommand += " --handover-timeout=" + handoverTimeout.String()
+	}
+
+	if profile != "" {
+		proxyCommand += " --profile=" + profile
+	}
+
+	return proxyCommand, nil
+}
+
 func generateHostConfig(opts SetupOptions) (string, error) {
 	identityFilePath, err := keys.GetLocalSSHKeyPath(opts.ClusterID, opts.SSHKeysDir)
 	if err != nil {
 		return "", fmt.Errorf("failed to get local keys folder: %w", err)
 	}
 
-	execPath, err := os.Executable()
+	proxyCommand, err := GenerateProxyCommand(opts.ClusterID, opts.AutoStartCluster, opts.ShutdownDelay, opts.Profile, "", 0, 0)
 	if err != nil {
-		return "", fmt.Errorf("failed to get executable path: %w", err)
-	}
-
-	profileOption := ""
-	if opts.Profile != "" {
-		profileOption = "--profile=" + opts.Profile
+		return "", fmt.Errorf("failed to generate ProxyCommand: %w", err)
 	}
 
 	hostConfig := fmt.Sprintf(`
@@ -77,8 +97,8 @@ Host %s
     ConnectTimeout 360
     StrictHostKeyChecking accept-new
     IdentityFile %q
-    ProxyCommand %q ssh connect --proxy --cluster=%s --auto-start-cluster=%t --shutdown-delay=%s %s
-`, opts.HostName, identityFilePath, execPath, opts.ClusterID, opts.AutoStartCluster, opts.ShutdownDelay, profileOption)
+    ProxyCommand %s
+`, opts.HostName, identityFilePath, proxyCommand)
 
 	return hostConfig, nil
 }

experimental/ssh/internal/setup/setup_test.go

@@ -54,6 +54,24 @@ func TestValidateClusterAccess_ClusterNotFound(t *testing.T) {
 	assert.Contains(t, err.Error(), "failed to get cluster information for cluster ID 'nonexistent'")
 }
 
+func TestGenerateProxyCommand(t *testing.T) {
+	cmd, err := GenerateProxyCommand("cluster-123", true, 45*time.Second, "", "", 0, 0)
+	assert.NoError(t, err)
+	assert.Contains(t, cmd, "ssh connect --proxy --cluster=cluster-123 --auto-start-cluster=true --shutdown-delay=45s")
+	assert.NotContains(t, cmd, "--metadata")
+	assert.NotContains(t, cmd, "--profile")
+	assert.NotContains(t, cmd, "--handover-timeout")
+}
+
+func TestGenerateProxyCommand_WithExtraArgs(t *testing.T) {
+	cmd, err := GenerateProxyCommand("cluster-123", true, 45*time.Second, "test-profile", "user", 2222, 2*time.Minute)
+	assert.NoError(t, err)
+	assert.Contains(t, cmd, "ssh connect --proxy --cluster=cluster-123 --auto-start-cluster=true --shutdown-delay=45s")
+	assert.Contains(t, cmd, " --metadata=user,2222")
+	assert.Contains(t, cmd, " --handover-timeout=2m0s")
+	assert.Contains(t, cmd, " --profile=test-profile")
+}
+
 func TestGenerateHostConfig_Valid(t *testing.T) {
 	// Create a temporary directory for testing
 	tmpDir := t.TempDir()