SSH: bump serverless environment version and expose a flag for it

experimental/ssh/cmd/connect.go

@@ -35,6 +35,7 @@ the SSH server and handling the connection proxy.
 	var userKnownHostsFile string
 	var liteswap string
 	var skipSettingsCheck bool
+	var environmentVersion int
 
 	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")
@@ -67,6 +68,9 @@ the SSH server and handling the connection proxy.
 	cmd.Flags().BoolVar(&skipSettingsCheck, "skip-settings-check", false, "Skip checking and updating IDE settings")
 	cmd.Flags().MarkHidden("skip-settings-check")
 
+	cmd.Flags().IntVar(&environmentVersion, "environment-version", defaultEnvironmentVersion, "Environment version for serverless compute")
+	cmd.Flags().MarkHidden("environment-version")
+
 	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 {
@@ -101,6 +105,7 @@ the SSH server and handling the connection proxy.
 			UserKnownHostsFile:   userKnownHostsFile,
 			Liteswap:             liteswap,
 			SkipSettingsCheck:    skipSettingsCheck,
+			EnvironmentVersion:   environmentVersion,
 			AdditionalArgs:       args,
 		}
 		if err := opts.Validate(); err != nil {

experimental/ssh/cmd/constants.go

@@ -3,10 +3,11 @@ package ssh
 import "time"
 
 const (
-	defaultServerPort      = 7772
-	defaultMaxClients      = 10
-	defaultShutdownDelay   = 10 * time.Minute
-	defaultHandoverTimeout = 30 * time.Minute
+	defaultServerPort         = 7772
+	defaultMaxClients         = 10
+	defaultShutdownDelay      = 10 * time.Minute
+	defaultHandoverTimeout    = 30 * time.Minute
+	defaultEnvironmentVersion = 4
 
 	serverTimeout        = 24 * time.Hour
 	taskStartupTimeout   = 10 * time.Minute

experimental/ssh/internal/client/client.go

@@ -44,6 +44,7 @@ var connectionNameRegex = regexp.MustCompile(`^[a-zA-Z0-9][a-zA-Z0-9_-]*$`)
 const (
 	sshServerTaskKey         = "start_ssh_server"
 	serverlessEnvironmentKey = "ssh_tunnel_serverless"
+	minEnvironmentVersion    = 4
 
 	VSCodeOption  = "vscode"
 	VSCodeCommand = "code"
@@ -98,6 +99,8 @@ type ClientOptions struct {
 	Liteswap string
 	// If true, skip checking and updating IDE settings.
 	SkipSettingsCheck bool
+	// Environment version for serverless compute.
+	EnvironmentVersion int
 }
 
 func (o *ClientOptions) Validate() error {
@@ -117,6 +120,9 @@ func (o *ClientOptions) Validate() error {
 	if o.IDE != "" && o.IDE != VSCodeOption && o.IDE != CursorOption {
 		return fmt.Errorf("invalid IDE value: %q, expected %q or %q", o.IDE, VSCodeOption, CursorOption)
 	}
+	if o.EnvironmentVersion > 0 && o.EnvironmentVersion < minEnvironmentVersion {
+		return fmt.Errorf("environment version must be >= %d, got %d", minEnvironmentVersion, o.EnvironmentVersion)
+	}
 	return nil
 }
 
@@ -182,6 +188,10 @@ func (o *ClientOptions) ToProxyCommand() (string, error) {
 		proxyCommand += " --liteswap=" + o.Liteswap
 	}
 
+	if o.EnvironmentVersion > 0 {
+		proxyCommand += " --environment-version=" + strconv.Itoa(o.EnvironmentVersion)
+	}
+
 	return proxyCommand, nil
 }
 
@@ -508,7 +518,7 @@ func submitSSHTunnelJob(ctx context.Context, client *databricks.WorkspaceClient,
 			{
 				EnvironmentKey: serverlessEnvironmentKey,
 				Spec: &compute.Environment{
-					EnvironmentVersion: "3",
+					EnvironmentVersion: strconv.Itoa(max(opts.EnvironmentVersion, minEnvironmentVersion)),
 				},
 			},
 		}

experimental/ssh/internal/client/client_test.go

@@ -76,6 +76,15 @@ func TestValidate(t *testing.T) {
 			name: "valid IDE cursor",
 			opts: client.ClientOptions{ClusterID: "abc-123", IDE: "cursor"},
 		},
+		{
+			name:    "environment version too low",
+			opts:    client.ClientOptions{ClusterID: "abc-123", EnvironmentVersion: 3},
+			wantErr: "environment version must be >= 4, got 3",
+		},
+		{
+			name: "valid environment version",
+			opts: client.ClientOptions{ClusterID: "abc-123", EnvironmentVersion: 4},
+		},
 	}
 
 	for _, tt := range tests {
@@ -140,6 +149,11 @@ func TestToProxyCommand(t *testing.T) {
 			opts: client.ClientOptions{ClusterID: "abc-123", Liteswap: "test-env"},
 			want: quoted + " ssh connect --proxy --cluster=abc-123 --auto-start-cluster=false --shutdown-delay=0s --liteswap=test-env",
 		},
+		{
+			name: "with environment version",
+			opts: client.ClientOptions{ClusterID: "abc-123", EnvironmentVersion: 4},
+			want: quoted + " ssh connect --proxy --cluster=abc-123 --auto-start-cluster=false --shutdown-delay=0s --environment-version=4",
+		},
 	}
 
 	for _, tt := range tests {