Address review comments: identity tracking, safe cleanup, no StrictHostKeyChecking=no

- Add UserName to Session struct and include in FindMatching to prevent
  cross-identity session mixups when switching profiles
- Only run cleanup on definitive errServerMetadata errors; log and skip
  on transient failures (network, auth) to avoid deleting live sessions
- Add workspace host hash to generated session names to avoid SSH
  known_hosts conflicts across workspaces, removing the need for
  StrictHostKeyChecking=no and UserKnownHostsFile=/dev/null
- Prune expired sessions from disk during FindMatching
- Make resolveServerlessSession a method on ClientOptions
- Handle rand.Read error explicitly

Co-authored-by: Isaac

Author: anton-107

experimental/ssh/internal/client/client.go

@@ -211,8 +211,7 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt
 
 	// For serverless without explicit --name: auto-generate or reconnect to existing session.
 	if opts.IsServerlessMode() && opts.ConnectionName == "" && !opts.ProxyMode {
-		err := resolveServerlessSession(ctx, client, &opts)
-		if err != nil {
+		if err := opts.resolveServerlessSession(ctx, client); err != nil {
 			return err
 		}
 	}
@@ -343,10 +342,16 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt
 
 	// Persist the session for future reconnects.
 	if opts.IsServerlessMode() && !opts.ProxyMode {
+		currentUser, userErr := client.CurrentUser.Me(ctx)
+		sessionUserName := ""
+		if userErr == nil {
+			sessionUserName = currentUser.UserName
+		}
 		err = sessions.Add(ctx, sessions.Session{
 			Name:          opts.ConnectionName,
 			Accelerator:   opts.Accelerator,
 			WorkspaceHost: client.Config.Host,
+			UserName:      sessionUserName,
 			CreatedAt:     time.Now(),
 			ClusterID:     clusterID,
 		})
@@ -407,12 +412,7 @@ func ensureSSHConfigEntry(ctx context.Context, configPath, hostName, userName, k
 		return fmt.Errorf("failed to generate ProxyCommand: %w", err)
 	}
 
-	var hostConfig string
-	if opts.IsServerlessMode() {
-		hostConfig = sshconfig.GenerateServerlessHostConfig(hostName, userName, keyPath, proxyCommand)
-	} else {
-		hostConfig = sshconfig.GenerateHostConfig(hostName, userName, keyPath, proxyCommand)
-	}
+	hostConfig := sshconfig.GenerateHostConfig(hostName, userName, keyPath, proxyCommand)
 
 	_, err = sshconfig.CreateOrUpdateHostConfig(ctx, hostName, hostConfig, true)
 	if err != nil {
@@ -580,22 +580,15 @@ func spawnSSHClient(ctx context.Context, userName, privateKeyPath string, server
 
 	hostName := opts.SessionIdentifier()
 
-	hostKeyChecking := "StrictHostKeyChecking=accept-new"
-	if opts.IsServerlessMode() {
-		hostKeyChecking = "StrictHostKeyChecking=no"
-	}
-
 	sshArgs := []string{
 		"-l", userName,
 		"-i", privateKeyPath,
 		"-o", "IdentitiesOnly=yes",
-		"-o", hostKeyChecking,
+		"-o", "StrictHostKeyChecking=accept-new",
 		"-o", "ConnectTimeout=360",
 		"-o", "ProxyCommand=" + proxyCommand,
 	}
-	if opts.IsServerlessMode() {
-		sshArgs = append(sshArgs, "-o", "UserKnownHostsFile=/dev/null")
-	} else if opts.UserKnownHostsFile != "" {
+	if opts.UserKnownHostsFile != "" {
 		sshArgs = append(sshArgs, "-o", "UserKnownHostsFile="+opts.UserKnownHostsFile)
 	}
 	sshArgs = append(sshArgs, hostName)
@@ -740,12 +733,17 @@ func ensureSSHServerIsRunning(ctx context.Context, client *databricks.WorkspaceC
 }
 
 // resolveServerlessSession handles auto-generation and reconnection for serverless sessions.
-// It checks local state for existing sessions matching the workspace and accelerator,
+// It checks local state for existing sessions matching the workspace, accelerator, and user,
 // probes them to see if they're still alive, and prompts the user to reconnect or create new.
-func resolveServerlessSession(ctx context.Context, client *databricks.WorkspaceClient, opts *ClientOptions) error {
+func (opts *ClientOptions) resolveServerlessSession(ctx context.Context, client *databricks.WorkspaceClient) error {
 	version := build.GetInfo().Version
 
-	matching, err := sessions.FindMatching(ctx, client.Config.Host, opts.Accelerator)
+	me, err := client.CurrentUser.Me(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to get current user: %w", err)
+	}
+
+	matching, err := sessions.FindMatching(ctx, client.Config.Host, opts.Accelerator, me.UserName)
 	if err != nil {
 		log.Warnf(ctx, "Failed to load session state: %v", err)
 	}
@@ -761,8 +759,12 @@ func resolveServerlessSession(ctx context.Context, client *databricks.WorkspaceC
 		_, _, _, probeErr := getServerMetadata(ctx, client, s.Name, s.ClusterID, version, opts.Liteswap)
 		if probeErr == nil {
 			alive = append(alive, s)
-		} else {
+		} else if errors.Is(probeErr, errServerMetadata) {
+			// Only clean up when the server is definitively gone (metadata endpoint returns not-found).
+			// Transient errors (network, auth) should not trigger cleanup.
 			cleanupStaleSession(ctx, client, s, version)
+		} else {
+			log.Warnf(ctx, "Transient error probing session %s, skipping: %v", s.Name, probeErr)
 		}
 	}
 
@@ -788,7 +790,7 @@ func resolveServerlessSession(ctx context.Context, client *databricks.WorkspaceC
 	}
 
 	// No alive session selected — generate a new name.
-	opts.ConnectionName = sessions.GenerateSessionName(opts.Accelerator)
+	opts.ConnectionName = sessions.GenerateSessionName(opts.Accelerator, client.Config.Host)
 	cmdio.LogString(ctx, "Creating new session: "+opts.ConnectionName)
 	return nil
 }

experimental/ssh/internal/sessions/namegen.go

@@ -1,8 +1,10 @@
 package sessions
 
 import (
+	"crypto/md5"
 	"crypto/rand"
 	"encoding/hex"
+	"fmt"
 	"strings"
 	"time"
 )
@@ -13,16 +15,26 @@ var acceleratorPrefixes = map[string]string{
 	"GPU_8xH100": "gpu-h100",
 }
 
-// GenerateSessionName creates a human-readable session name from the accelerator type.
-// Format: <prefix>-<random_hex>, e.g. "gpu-a10-f3a2b1c0".
-func GenerateSessionName(accelerator string) string {
+// GenerateSessionName creates a human-readable session name from the accelerator type
+// and workspace host. The workspace host is hashed into the name to avoid SSH known_hosts
+// conflicts when connecting to different workspaces.
+// Format: databricks-<prefix>-<date>-<workspace_hash><random_hex>.
+func GenerateSessionName(accelerator, workspaceHost string) string {
 	prefix, ok := acceleratorPrefixes[accelerator]
 	if !ok {
 		prefix = strings.ToLower(strings.ReplaceAll(accelerator, "_", "-"))
 	}
 
 	date := time.Now().Format("20060102")
+
+	// Include a short hash of the workspace host to avoid known_hosts conflicts
+	// when connecting to different workspaces.
+	wsHash := md5.Sum([]byte(workspaceHost))
+	wsHashStr := hex.EncodeToString(wsHash[:])[:4]
+
 	b := make([]byte, 3)
-	_, _ = rand.Read(b)
-	return "databricks-" + prefix + "-" + date + "-" + hex.EncodeToString(b)
+	if _, err := rand.Read(b); err != nil {
+		panic(fmt.Sprintf("crypto/rand.Read failed: %v", err))
+	}
+	return "databricks-" + prefix + "-" + date + "-" + wsHashStr + hex.EncodeToString(b)
 }

experimental/ssh/internal/sessions/sessions.go

@@ -23,6 +23,7 @@ type Session struct {
 	Name          string    `json:"name"`
 	Accelerator   string    `json:"accelerator"`
 	WorkspaceHost string    `json:"workspace_host"`
+	UserName      string    `json:"user_name,omitempty"`
 	CreatedAt     time.Time `json:"created_at"`
 	ClusterID     string    `json:"cluster_id,omitempty"`
 }
@@ -129,17 +130,35 @@ func Remove(ctx context.Context, name string) error {
 	return Save(ctx, store)
 }
 
-// FindMatching returns non-expired sessions that match the given workspace host and accelerator.
-func FindMatching(ctx context.Context, workspaceHost, accelerator string) ([]Session, error) {
+// FindMatching returns non-expired sessions that match the given workspace host, accelerator,
+// and user name. Expired sessions are pruned from the store on disk.
+func FindMatching(ctx context.Context, workspaceHost, accelerator, userName string) ([]Session, error) {
 	store, err := Load(ctx)
 	if err != nil {
 		return nil, err
 	}
 
 	cutoff := time.Now().Add(-sessionMaxAge)
-	var result []Session
+
+	// Prune expired sessions from the store.
+	active := store.Sessions[:0]
+	pruned := false
 	for _, s := range store.Sessions {
-		if s.WorkspaceHost == workspaceHost && s.Accelerator == accelerator && s.CreatedAt.After(cutoff) {
+		if s.CreatedAt.After(cutoff) {
+			active = append(active, s)
+		} else {
+			pruned = true
+		}
+	}
+	if pruned {
+		store.Sessions = active
+		// Best-effort save; don't fail the operation if pruning fails.
+		_ = Save(ctx, store)
+	}
+
+	var result []Session
+	for _, s := range active {
+		if s.WorkspaceHost == workspaceHost && s.Accelerator == accelerator && s.UserName == userName {
 			result = append(result, s)
 		}
 	}

experimental/ssh/internal/sessions/sessions_test.go

@@ -90,32 +90,41 @@ func TestFindMatching(t *testing.T) {
 
 	ctx := t.Context()
 	host := "https://test.databricks.com"
+	user := "alice@example.com"
 
 	now := time.Now()
 
-	err := Add(ctx, Session{Name: "s1", Accelerator: "GPU_1xA10", WorkspaceHost: host, CreatedAt: now})
+	err := Add(ctx, Session{Name: "s1", Accelerator: "GPU_1xA10", WorkspaceHost: host, UserName: user, CreatedAt: now})
 	require.NoError(t, err)
-	err = Add(ctx, Session{Name: "s2", Accelerator: "GPU_8xH100", WorkspaceHost: host, CreatedAt: now})
+	err = Add(ctx, Session{Name: "s2", Accelerator: "GPU_8xH100", WorkspaceHost: host, UserName: user, CreatedAt: now})
 	require.NoError(t, err)
-	err = Add(ctx, Session{Name: "s3", Accelerator: "GPU_1xA10", WorkspaceHost: "https://other.com", CreatedAt: now})
+	err = Add(ctx, Session{Name: "s3", Accelerator: "GPU_1xA10", WorkspaceHost: "https://other.com", UserName: user, CreatedAt: now})
 	require.NoError(t, err)
-	err = Add(ctx, Session{Name: "s4", Accelerator: "GPU_1xA10", WorkspaceHost: host, CreatedAt: now})
+	err = Add(ctx, Session{Name: "s4", Accelerator: "GPU_1xA10", WorkspaceHost: host, UserName: user, CreatedAt: now})
+	require.NoError(t, err)
+	err = Add(ctx, Session{Name: "s5", Accelerator: "GPU_1xA10", WorkspaceHost: host, UserName: "bob@example.com", CreatedAt: now})
 	require.NoError(t, err)
 
-	matches, err := FindMatching(ctx, host, "GPU_1xA10")
+	matches, err := FindMatching(ctx, host, "GPU_1xA10", user)
 	require.NoError(t, err)
 	assert.Len(t, matches, 2)
 	assert.Equal(t, "s1", matches[0].Name)
 	assert.Equal(t, "s4", matches[1].Name)
 
-	matches, err = FindMatching(ctx, host, "GPU_8xH100")
+	matches, err = FindMatching(ctx, host, "GPU_8xH100", user)
 	require.NoError(t, err)
 	assert.Len(t, matches, 1)
 	assert.Equal(t, "s2", matches[0].Name)
 
-	matches, err = FindMatching(ctx, host, "GPU_4xA100")
+	matches, err = FindMatching(ctx, host, "GPU_4xA100", user)
 	require.NoError(t, err)
 	assert.Empty(t, matches)
+
+	// Different user should not see alice's sessions
+	matches, err = FindMatching(ctx, host, "GPU_1xA10", "bob@example.com")
+	require.NoError(t, err)
+	assert.Len(t, matches, 1)
+	assert.Equal(t, "s5", matches[0].Name)
 }
 
 func TestFindMatchingExpiresOldSessions(t *testing.T) {
@@ -125,16 +134,22 @@ func TestFindMatchingExpiresOldSessions(t *testing.T) {
 
 	ctx := t.Context()
 	host := "https://test.databricks.com"
+	user := "alice@example.com"
 
-	err := Add(ctx, Session{Name: "old", Accelerator: "GPU_1xA10", WorkspaceHost: host, CreatedAt: time.Now().Add(-25 * time.Hour)})
+	err := Add(ctx, Session{Name: "old", Accelerator: "GPU_1xA10", WorkspaceHost: host, UserName: user, CreatedAt: time.Now().Add(-25 * time.Hour)})
 	require.NoError(t, err)
-	err = Add(ctx, Session{Name: "recent", Accelerator: "GPU_1xA10", WorkspaceHost: host, CreatedAt: time.Now()})
+	err = Add(ctx, Session{Name: "recent", Accelerator: "GPU_1xA10", WorkspaceHost: host, UserName: user, CreatedAt: time.Now()})
 	require.NoError(t, err)
 
-	matches, err := FindMatching(ctx, host, "GPU_1xA10")
+	matches, err := FindMatching(ctx, host, "GPU_1xA10", user)
 	require.NoError(t, err)
 	require.Len(t, matches, 1)
 	assert.Equal(t, "recent", matches[0].Name)
+
+	// Verify expired sessions were pruned from disk.
+	store, err := Load(ctx)
+	require.NoError(t, err)
+	assert.Len(t, store.Sessions, 1, "expired sessions should be pruned from disk")
 }
 
 func TestStateFilePath(t *testing.T) {
@@ -151,6 +166,7 @@ func TestStateFilePath(t *testing.T) {
 var connectionNameRegex = regexp.MustCompile(`^[a-zA-Z0-9][a-zA-Z0-9_-]*$`)
 
 func TestGenerateSessionName(t *testing.T) {
+	host := "https://test.databricks.com"
 	tests := []struct {
 		accelerator    string
 		wantPrefix     string
@@ -163,7 +179,7 @@ func TestGenerateSessionName(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.accelerator, func(t *testing.T) {
-			name := GenerateSessionName(tt.accelerator)
+			name := GenerateSessionName(tt.accelerator, host)
 			assert.Greater(t, len(name), len(tt.wantPrefix), "name should be longer than prefix")
 			assert.Equal(t, tt.wantPrefix, name[:len(tt.wantPrefix)])
 			// Verify date component is present (starts with "20" for 2000s dates).
@@ -173,10 +189,20 @@ func TestGenerateSessionName(t *testing.T) {
 	}
 }
 
+func TestGenerateSessionNameDiffersByWorkspace(t *testing.T) {
+	name1 := GenerateSessionName("GPU_1xA10", "https://workspace-a.databricks.com")
+	name2 := GenerateSessionName("GPU_1xA10", "https://workspace-b.databricks.com")
+	// The workspace hash portion (after the date-) should differ.
+	// Names have format: databricks-gpu-a10-YYYYMMDD-<wshash><random>
+	// Extract after the date prefix to compare workspace hash parts.
+	assert.NotEqual(t, name1, name2, "names for different workspaces should differ")
+}
+
 func TestGenerateSessionNameUniqueness(t *testing.T) {
+	host := "https://test.databricks.com"
 	seen := make(map[string]bool)
 	for range 100 {
-		name := GenerateSessionName("GPU_1xA10")
+		name := GenerateSessionName("GPU_1xA10", host)
 		assert.False(t, seen[name], "duplicate name generated: %s", name)
 		seen[name] = true
 	}

experimental/ssh/internal/sshconfig/sshconfig.go

@@ -175,30 +175,13 @@ func RemoveHostConfig(ctx context.Context, hostName string) error {
 
 // GenerateHostConfig generates an SSH host config block.
 func GenerateHostConfig(hostName, userName, identityFile, proxyCommand string) string {
-	return generateHostConfig(hostName, userName, identityFile, proxyCommand, false)
-}
-
-// GenerateServerlessHostConfig generates an SSH host config block for serverless compute.
-// It disables strict host key checking since serverless containers generate fresh keys each time,
-// and identity is already verified through Databricks authentication and Driver Proxy.
-func GenerateServerlessHostConfig(hostName, userName, identityFile, proxyCommand string) string {
-	return generateHostConfig(hostName, userName, identityFile, proxyCommand, true)
-}
-
-func generateHostConfig(hostName, userName, identityFile, proxyCommand string, serverless bool) string {
-	hostKeyChecking := "StrictHostKeyChecking accept-new"
-	knownHostsLine := ""
-	if serverless {
-		hostKeyChecking = "StrictHostKeyChecking no"
-		knownHostsLine = "    UserKnownHostsFile /dev/null\n"
-	}
 	return fmt.Sprintf(`
 Host %s
     User %s
     ConnectTimeout 360
-    %s
-%s    IdentitiesOnly yes
+    StrictHostKeyChecking accept-new
+    IdentitiesOnly yes
     IdentityFile %q
     ProxyCommand %s
-`, hostName, userName, hostKeyChecking, knownHostsLine, identityFile, proxyCommand)
+`, hostName, userName, identityFile, proxyCommand)
 }