Auto update IDE settings for serverless ssh mode

experimental/ssh/cmd/connect.go

@@ -35,6 +35,7 @@ the SSH server and handling the connection proxy.
 	var autoStartCluster bool
 	var userKnownHostsFile string
 	var liteswap string
+	var skipSettingsCheck bool
 
 	cmd.Flags().StringVar(&clusterID, "cluster", "", "Databricks cluster ID (for dedicated clusters)")
 	cmd.Flags().DurationVar(&shutdownDelay, "shutdown-delay", defaultShutdownDelay, "Delay before shutting down the server after the last client disconnects")
@@ -64,6 +65,9 @@ the SSH server and handling the connection proxy.
 	cmd.Flags().StringVar(&liteswap, "liteswap", "", "Liteswap header value for traffic routing (dev/test only)")
 	cmd.Flags().MarkHidden("liteswap")
 
+	cmd.Flags().BoolVar(&skipSettingsCheck, "skip-settings-check", false, "Skip checking and updating IDE settings")
+	cmd.Flags().MarkHidden("skip-settings-check")
+
 	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 {
@@ -113,6 +117,7 @@ the SSH server and handling the connection proxy.
 			ClientPrivateKeyName: clientPrivateKeyName,
 			UserKnownHostsFile:   userKnownHostsFile,
 			Liteswap:             liteswap,
+			SkipSettingsCheck:    skipSettingsCheck,
 			AdditionalArgs:       args,
 		}
 		return client.Run(ctx, wsClient, opts)

experimental/ssh/internal/client/client.go

@@ -20,6 +20,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/sshconfig"
+	"github.com/databricks/cli/experimental/ssh/internal/vscode"
 	sshWorkspace "github.com/databricks/cli/experimental/ssh/internal/workspace"
 	"github.com/databricks/cli/internal/build"
 	"github.com/databricks/cli/libs/cmdio"
@@ -92,6 +93,8 @@ type ClientOptions struct {
 	UserKnownHostsFile string
 	// Liteswap header value for traffic routing (dev/test only).
 	Liteswap string
+	// If true, skip checking and updating IDE settings.
+	SkipSettingsCheck bool
 }
 
 func (o *ClientOptions) IsServerlessMode() bool {
@@ -206,6 +209,26 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt
 	cmdio.LogString(ctx, "Using SSH key: "+keyPath)
 	cmdio.LogString(ctx, fmt.Sprintf("Secrets scope: %s, key name: %s", secretScopeName, opts.ClientPublicKeyName))
 
+	// Check and update IDE settings for serverless mode, where we must set up
+	// desired server ports (or socket connection mode) for the connection to go through
+	// (as the majority of the localhost ports on the remote side are blocked by iptable rules).
+	// Plus the platform (always linux), and extensions (python and jupyter), to make the initial experience smoother.
+	if opts.IDE != "" && opts.IsServerlessMode() && !opts.ProxyMode && !opts.SkipSettingsCheck && cmdio.IsPromptSupported(ctx) {
+		err = vscode.CheckAndUpdateSettings(ctx, opts.IDE, opts.ConnectionName)
+		if err != nil {
+			cmdio.LogString(ctx, fmt.Sprintf("Failed to update IDE settings: %v", err))
+			cmdio.LogString(ctx, vscode.GetManualInstructions(opts.IDE, opts.ConnectionName))
+			cmdio.LogString(ctx, "Use --skip-settings-check to bypass IDE settings verification.")
+			shouldProceed, promptErr := cmdio.AskYesOrNo(ctx, "Do you want to proceed with the connection?")
+			if promptErr != nil {
+				return fmt.Errorf("failed to prompt user: %w", promptErr)
+			}
+			if !shouldProceed {
+				return errors.New("aborted: IDE settings need to be updated manually, user declined to proceed")
+			}
+		}
+	}
+
 	var userName string
 	var serverPort int
 	var clusterID string