ssh: document --environment-version on ssh connect

.nextchanges/cli/ssh-environment-version.md

@@ -0,0 +1 @@
+* `databricks ssh connect` now documents `--environment-version` for choosing the serverless environment version, starts a separate default-named server for each version, and rejects the flag together with `--cluster`. ([#7264](https://github.com/databricks/cli/pull/7264))

experimental/ssh/cmd/connect.go

@@ -36,6 +36,7 @@ Connect to serverless:
   databricks ssh connect
   databricks ssh connect --accelerator=<GPU_type>   # AI Runtime
   databricks ssh connect --base-environment=<name>  # custom base environment
+  databricks ssh connect --environment-version=5    # serverless environment version
 
 Connect to a dedicated cluster:
   databricks ssh connect --cluster=<cluster-id>`,
@@ -93,8 +94,7 @@ Connect to a dedicated cluster:
 	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 AI Runtime")
-	cmd.Flags().MarkHidden("environment-version")
+	cmd.Flags().IntVar(&environmentVersion, "environment-version", defaultEnvironmentVersion, "Serverless environment version for the SSH server job (serverless only)")
 
 	cmd.Flags().StringVar(&baseEnvironment, "base-environment", "", "Custom base environment for serverless compute: an env.yaml path, a workspace-base-environments resource ID, or a display name")
 
@@ -114,21 +114,21 @@ Connect to a dedicated cluster:
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
 		ctx := cmd.Context()
 		wsClient := cmdctx.WorkspaceClient(ctx)
+		// Only carry an explicitly-set environment version. Leaving it at 0 otherwise
+		// lets the submit path default to minEnvironmentVersion and lets Validate
+		// detect a real --environment-version + --base-environment conflict.
+		if !cmd.Flags().Changed("environment-version") {
+			environmentVersion = 0
+		}
 		if connectionName == "" && clusterID == "" && !proxyMode {
-			connectionName = client.GenerateDefaultConnectionName(wsClient.Config.Host, accelerator, baseEnvironment)
+			connectionName = client.GenerateDefaultConnectionName(wsClient.Config.Host, accelerator, baseEnvironment, environmentVersion)
 		}
 		// Serverless GPU compute can take much longer to provision than CPU compute,
 		// so allow extra time for the SSH server job to start.
 		startupTimeout := taskStartupTimeout
 		if accelerator != "" {
 			startupTimeout = gpuTaskStartupTimeout
 		}
-		// Only carry an explicitly-set environment version. Leaving it at 0 otherwise
-		// lets the submit path default to minEnvironmentVersion and lets Validate
-		// detect a real --environment-version + --base-environment conflict.
-		if !cmd.Flags().Changed("environment-version") {
-			environmentVersion = 0
-		}
 		opts := client.ClientOptions{
 			Profile:               wsClient.Config.Profile,
 			ClusterID:             clusterID,

experimental/ssh/internal/client/client.go

@@ -181,6 +181,9 @@ func (o *ClientOptions) Validate() error {
 	if o.BaseEnvironment != "" && o.ClusterID != "" {
 		return errors.New("--base-environment can only be used with serverless compute")
 	}
+	if o.EnvironmentVersion > 0 && o.ClusterID != "" {
+		return errors.New("--environment-version can only be used with serverless compute")
+	}
 	// A server started with fewer than one client slot rejects every connection with an
 	// opaque websocket handshake failure, so catch it here instead.
 	if o.MaxClients < 1 {
@@ -201,18 +204,23 @@ func (o *ClientOptions) Validate() error {
 }
 
 // GenerateDefaultConnectionName creates a deterministic connection name from
-// the workspace host, accelerator type, and base environment. The name includes
+// the workspace host, accelerator type, and environment. The name includes
 // a hash so that different workspaces produce different names (avoiding SSH
 // known_hosts conflicts). The environment is folded into the hash because a
 // serverless server bakes in its environment at startup: distinct environments
 // must map to distinct connection names so they don't reuse each other's server.
-func GenerateDefaultConnectionName(host, accelerator, baseEnvironment string) string {
-	// Keep the hash host-only when no base environment is set so existing default
+func GenerateDefaultConnectionName(host, accelerator, baseEnvironment string, environmentVersion int) string {
+	// Keep the hash host-only when no environment is set so existing default
 	// connection names are preserved.
 	hashInput := host
 	if baseEnvironment != "" {
 		hashInput = host + "\x00" + baseEnvironment
 	}
+	// The two are mutually exclusive. The key keeps a version from hashing like a
+	// base environment whose name is that number.
+	if environmentVersion > 0 {
+		hashInput = host + "\x00environment_version=" + strconv.Itoa(environmentVersion)
+	}
 	h := md5.Sum([]byte(hashInput))
 	hashStr := hex.EncodeToString(h[:4])
 	if accelerator != "" {

experimental/ssh/internal/client/client_test.go

@@ -92,7 +92,12 @@ func TestValidate(t *testing.T) {
 		},
 		{
 			name: "valid environment version",
-			opts: client.ClientOptions{ClusterID: "abc-123", EnvironmentVersion: 4},
+			opts: client.ClientOptions{ConnectionName: "my-conn", EnvironmentVersion: 4},
+		},
+		{
+			name:    "environment version with cluster",
+			opts:    client.ClientOptions{ClusterID: "abc-123", EnvironmentVersion: 5},
+			wantErr: "--environment-version can only be used with serverless compute",
 		},
 		{
 			name:    "base environment with environment version",
@@ -256,7 +261,7 @@ func TestGenerateDefaultConnectionName(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			got := client.GenerateDefaultConnectionName(tt.host, tt.accelerator, "")
+			got := client.GenerateDefaultConnectionName(tt.host, tt.accelerator, "", 0)
 			assert.Equal(t, tt.want, got)
 		})
 	}
@@ -266,12 +271,25 @@ func TestGenerateDefaultConnectionName(t *testing.T) {
 	// ignored when an existing server for the default name is reused).
 	t.Run("base environment differentiates the name", func(t *testing.T) {
 		const host = "https://my-workspace.cloud.databricks.com"
-		base := client.GenerateDefaultConnectionName(host, "", "")
-		withEnv := client.GenerateDefaultConnectionName(host, "", "my-env")
-		otherEnv := client.GenerateDefaultConnectionName(host, "", "other-env")
+		base := client.GenerateDefaultConnectionName(host, "", "", 0)
+		withEnv := client.GenerateDefaultConnectionName(host, "", "my-env", 0)
+		otherEnv := client.GenerateDefaultConnectionName(host, "", "other-env", 0)
 		assert.NotEqual(t, base, withEnv, "setting --base-environment must change the default name")
 		assert.NotEqual(t, withEnv, otherEnv, "different environments must produce different names")
-		assert.Equal(t, withEnv, client.GenerateDefaultConnectionName(host, "", "my-env"), "must be deterministic")
+		assert.Equal(t, withEnv, client.GenerateDefaultConnectionName(host, "", "my-env", 0), "must be deterministic")
+	})
+
+	// Same for an explicit environment version: without it, a connection asking for
+	// version 5 would reuse a running server started on the default version.
+	t.Run("environment version differentiates the name", func(t *testing.T) {
+		const host = "https://my-workspace.cloud.databricks.com"
+		base := client.GenerateDefaultConnectionName(host, "", "", 0)
+		v5 := client.GenerateDefaultConnectionName(host, "", "", 5)
+		v6 := client.GenerateDefaultConnectionName(host, "", "", 6)
+		assert.NotEqual(t, base, v5, "setting --environment-version must change the default name")
+		assert.NotEqual(t, v5, v6, "different versions must produce different names")
+		assert.NotEqual(t, v5, client.GenerateDefaultConnectionName(host, "", "5", 0), "a version must not collide with a base environment named like it")
+		assert.Equal(t, v5, client.GenerateDefaultConnectionName(host, "", "", 5), "must be deterministic")
 	})
 }
 
@@ -288,7 +306,7 @@ func TestGenerateDefaultConnectionNameMatchesRegex(t *testing.T) {
 	for _, host := range hosts {
 		for _, acc := range accelerators {
 			for _, env := range environments {
-				name := client.GenerateDefaultConnectionName(host, acc, env)
+				name := client.GenerateDefaultConnectionName(host, acc, env, 0)
 				assert.Regexp(t, nameRegex, name, "host=%q accelerator=%q environment=%q name=%q", host, acc, env, name)
 			}
 		}