Auto-generate session names for serverless SSH connect#4701
Auto-generate session names for serverless SSH connect#4701anton-107 wants to merge 2 commits intossh-connect-elapsed-timefrom
Conversation
|
Commit: a177bea
16 interesting tests: 7 SKIP, 6 RECOVERED, 3 flaky
Top 20 slowest tests (at least 2 minutes):
|
…upport Remove the requirement for --name in serverless SSH connect. Sessions are now auto-generated with human-readable names (e.g. databricks-gpu-a10-20260310-a1b2c3), tracked in ~/.databricks/ssh-tunnel-sessions.json, and offered for reconnection on subsequent runs. Stale sessions are cleaned up automatically. Sessions expire after 24 hours. Also fixes known_hosts key mismatches for serverless by disabling strict host key checking (identity verified via Databricks auth). Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
d678d8d to
857404f
Compare
simonfaltum
left a comment
There was a problem hiding this comment.
[Agent Swarm Review] Verdict: Not ready yet
- 1 Critical
- 3 Major
- 1 Gap
- 2 Nit
- 2 Suggestion
The feature design (auto-generate session names, track sessions, offer reconnection) is good UX. However, there are several correctness issues that need addressing before this is safe to ship:
- Critical: Any probe failure triggers irreversible remote resource cleanup, even for transient network errors
- Major: No identity tracking means profile switches can cause cross-identity destructive cleanup
- Major: Expired sessions accumulate forever (remote resources leaked)
- Major: CLI version changes break reconnection within the 24h window
See inline comments for details.
| alive = append(alive, s) | ||
| } else { | ||
| cleanupStaleSession(ctx, client, s, version) | ||
| } |
There was a problem hiding this comment.
[Agent Swarm Review] [Critical]
Any probe error is treated as proof that the session is stale.
resolveServerlessSession() calls cleanupStaleSession() for every getServerMetadata() failure. That probe can fail for transient auth, network, workspace API, or version-mismatch reasons. In those cases the CLI will delete local SSH config, remove the session from state, and best-effort delete secret scopes and workspace content for a session that may still be alive.
Both reviewers flagged this. Isaac confirmed Critical in cross-review due to irreversible blast radius.
Suggestion: Only run destructive cleanup on definitive stale signals (e.g., 404/not-found). For transient errors, keep the session and surface a warning.
| Accelerator string `json:"accelerator"` | ||
| WorkspaceHost string `json:"workspace_host"` | ||
| CreatedAt time.Time `json:"created_at"` | ||
| ClusterID string `json:"cluster_id,omitempty"` |
There was a problem hiding this comment.
[Agent Swarm Review] [Major]
Session matching ignores Databricks identity. No username/profile is stored, so switching profiles on the same workspace causes cross-identity session mixup. Combined with the probe-failure cleanup issue, probing another identity's session fails and triggers cleanup of their remote resources.
Suggestion: Persist an identity key (e.g., workspace username) in the Session struct and include it in FindMatching.
| result = append(result, s) | ||
| } | ||
| } | ||
| return result, nil |
There was a problem hiding this comment.
[Agent Swarm Review] [Major]
Expired sessions accumulate in state file indefinitely. FindMatching filters out expired sessions from results but never removes them from the store on disk or triggers cleanup of their remote resources (secret scopes, workspace content). Over time the state file grows unboundedly and remote resources are leaked.
Suggestion: Prune expired sessions during Load or FindMatching, saving the cleaned store back.
|
|
||
| date := time.Now().Format("20060102") | ||
| b := make([]byte, 3) | ||
| _, _ = rand.Read(b) |
There was a problem hiding this comment.
[Agent Swarm Review] [Nit]
_, _ = rand.Read(b) discards the error. If crypto/rand.Read fails, b is all zeros, producing predictable session names. Consider checking the error or adding a comment explaining why it's intentionally ignored.
| hostKeyChecking := "StrictHostKeyChecking=accept-new" | ||
| if opts.IsServerlessMode() { | ||
| hostKeyChecking = "StrictHostKeyChecking=no" | ||
| } | ||
|
|
||
| sshArgs := []string{ | ||
| "-l", userName, | ||
| "-i", privateKeyPath, | ||
| "-o", "IdentitiesOnly=yes", | ||
| "-o", "StrictHostKeyChecking=accept-new", | ||
| "-o", hostKeyChecking, | ||
| "-o", "ConnectTimeout=360", | ||
| "-o", "ProxyCommand=" + proxyCommand, | ||
| } | ||
| if opts.UserKnownHostsFile != "" { | ||
| if opts.IsServerlessMode() { | ||
| sshArgs = append(sshArgs, "-o", "UserKnownHostsFile=/dev/null") | ||
| } else if opts.UserKnownHostsFile != "" { |
There was a problem hiding this comment.
We've had such options before, but the security didn't like it.
With auto host name generation we should not have that many host conflicts, right?
Before you would get them if you re-used the same name to connect to a different workspace. Re-using the same name for the same workspace is fine, as our server will get the server ssh key from the secrets scope that's tied to the name (and with the same name the scope will be the same). But across different workspaces we will get a problem, since server keys will be different.
Can we also add workspace id (real one, or based on the host url) to the generated session name?
| // resolveServerlessSession handles auto-generation and reconnection for serverless sessions. | ||
| // It checks local state for existing sessions matching the workspace and accelerator, | ||
| // 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 { |
There was a problem hiding this comment.
Nit, but this can be a method on the ClientOptions struct, might be easier to understand that we are mutating the options here then
| func resolveServerlessSession(ctx context.Context, client *databricks.WorkspaceClient, opts *ClientOptions) error { | ||
| version := build.GetInfo().Version | ||
|
|
||
| matching, err := sessions.FindMatching(ctx, client.Config.Host, opts.Accelerator) |
There was a problem hiding this comment.
I feel like majority of this logic can be moved to the sessions package (up until line 788). getServerMetadata can be passed as an argument. Then it will be easier to test.
Same for cleanupStaleSession. Or will there be circular dependencies it we do that? (since that function has a lot of them)
|
|
||
| // 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 { |
There was a problem hiding this comment.
As mentioned above, will it help with known_hosts conflicts if we add a workspace id/host here?
Summary
--namein serverlessssh connect;--acceleratoralone is now sufficientdatabricks-gpu-a10-20260310-a1b2c3)~/.databricks/ssh-tunnel-sessions.jsonand offer reconnection on subsequent runsStrictHostKeyChecking=no+UserKnownHostsFile=/dev/nullStacked on #4697.
Test plan
databricks ssh connect --accelerator GPU_1xA10creates new session🤖 Generated with Claude Code