sandbox: drive SSH config from server-stamped gateway_host (drop heuristic, --gateway flag, cursor-hiding spinner, add multi-gateway support) (#5543)

## Summary

Consolidates the gateway-resolution path simplification into a single
PR. Now that the server stamps `gateway_host` on every SshKey and
Sandbox response, `register` always populates the cache, every
subsequent operation that touches the API gets the workspace's real
gateway directly, and a bunch of compensating code falls out as dead
weight.

## What changes

### Code that goes away

- **`resolveGatewayHost(workspaceHost string) string`** — the hardcoded
heuristic that returned `uw2.dbrx.dev` for non-staging and
`ue1.s.dbrx.dev` for staging. Wrong for any workspace whose gateway
doesn't match those two values.
- **`defaultGatewayHost` + `stagingDefaultGatewayHost`** — the two
constants behind the heuristic.
- **`--gateway` flag on `sandbox ssh`** — a manual override that made
sense only when the CLI couldn't learn the gateway itself. With the
cache always populated post-register, the only thing the flag does now
is let a user typo their way into dialing a random host.
- **The "Connecting…" Bubble Tea spinner before `execSSHDirect`** — it
hid the terminal cursor at init (`initTerminal()`) and the deferred
`Close()` never ran because `execv` replaces the CLI process. Result:
cursor invisible for the entire ssh session and after it exited.
Co-authored with Anwell — pulled in from #5545.
- **Synthetic `Host sandbox-gw` alias** — UI deep links use the literal
gateway hostname, so the alias was buying nothing and adding an
indirection the UI had to know about.

### New behavior

- **Server-stamped gateway flows through to the SSH config.**
`registerKey` now returns `*sshKeyEntry` carrying `GatewayHost`;
`register` caches it via `setGatewayHost` and passes it to
`maybeWriteSSHConfig`.
- **One Host stanza per cached gateway.** A user with profiles in
multiple regions accumulates one block per unique gateway, instead of
the second `register` overwriting the first's block. New
`allGatewayHosts` helper in `state.go` returns the deduplicated set.
- **`saveState` is atomic** (tmp + Rename) — concurrent CLI invocations
can't see a half-written `sandbox.json` and trip a parse error that
`loadState` silently swallows.
- **Post-register hint branches on `getDefault`** — suggests `databricks
sandbox create` when no default exists, `databricks sandbox ssh` when
one does.
- **Path output uses forward slashes on Windows** — `filepath.ToSlash`
on the `Generated SSH key at …` line so acceptance goldens are stable
across Linux/macOS/Windows.

### New SSH config block shape

Before (9 lines, synthetic alias, several overrides):

```
# Managed by `databricks sandbox register`.
Host sandbox-gw
    HostName uw2.dbrx.dev
    Port 2222
    IdentityFile ~/.ssh/sandbox_ed25519
    IdentitiesOnly yes
    StrictHostKeyChecking no
    UserKnownHostsFile /dev/null
    LogLevel ERROR
```

After (literal hostname, one stanza per gateway, only the directives
that actually differ from SSH defaults):

```
# Managed by `databricks sandbox register`. Manual edits will be overwritten.
Host us-west-2.service-direct.dev.databricks.com
    Port 2222
    IdentityFile ~/.ssh/sandbox_ed25519
    IdentitiesOnly yes
Host us-east-1.service-direct.prod.databricks.com
    Port 2222
    IdentityFile ~/.ssh/sandbox_ed25519
    IdentitiesOnly yes
```

Mirrors the workspace UI's "First time setup?" snippet exactly, so CLI
users and paste-the-snippet users converge.

### New resolution chain in `sandbox ssh`

Fresh API response → cached value. If both are empty:

```
Error: could not resolve the sandbox SSH gateway — run `databricks sandbox register` first
```

## Tests

- `acceptance/cmd/sandbox/register/success/` — pins the new contract:
registerKey returns SshKey with gateway_host, register caches it, prints
the consent-skip notice in the non-interactive acceptance path. Parent
`test.toml` Ignores `.ssh`.
- `TestMaybeWriteSSHConfigSkipsWhenNoGatewaysCached` — pins the
no-gateway short-circuit.
- `TestBuildSSHConfigBlockMultipleGateways` — pins per-gateway
repetition in the new block shape.

## Test plan

- [x] `go test ./cmd/sandbox/...` passes
- [x] `go test ./acceptance -run TestAccept/cmd/sandbox` passes
- [x] `go build ./...` clean
- [x] `./task lint` clean
- [x] Windows variant passes (after `filepath.ToSlash` fix)

This pull request and its description were written by Isaac.

---------

Co-authored-by: Pieter Noordhuis <pieter.noordhuis@databricks.com>

Author: akshaysingla-db

acceptance/cmd/sandbox/register/success/out.test.toml

@@ -0,0 +1,9 @@
+  ✓ Generated SSH key at [TEST_TMP_DIR]/.ssh/sandbox_ed25519
+  ✓ SSH key registered
+  skipped SSH config update (non-interactive); re-run `databricks sandbox register` from a terminal to add the sandbox gateway block(s)
+
+  Run databricks sandbox create to provision your first sandbox.
+
+{
+  "[DATABRICKS_URL]": "us-west-2.service-direct.dev.databricks.com"
+}

acceptance/cmd/sandbox/register/success/output.txt

@@ -0,0 +1,6 @@
+$CLI sandbox register --name test-machine
+
+# Verify the gateway host was cached in the local state file.
+# The framework Ignores the .databricks directory by default, so we
+# print the cached value to stdout for assertion instead.
+cat "$HOME/.databricks/sandbox.json" | jq -r '.gatewayHosts // {}'

acceptance/cmd/sandbox/register/success/script

@@ -0,0 +1,17 @@
+[[Server]]
+Pattern = "POST /api/2.0/lakebox/ssh-keys"
+Response.Body = '''
+{
+  "keyHash": "a0b1c2d3e4f5a0b1c2d3e4f5a0b1c2d3",
+  "name": "test-machine",
+  "createTime": "2026-06-10T17:00:00Z",
+  "gatewayHost": "us-west-2.service-direct.dev.databricks.com"
+}
+'''
+
+# Pin the synthetic test hash so the framework's generic 3+-digit-run →
+# [NUMID] regex doesn't mangle it.
+[[Repls]]
+Old = 'a0b1c2d3e4f5a0b1c2d3e4f5a0b1c2d3'
+New = '[KEY_HASH]'
+Order = 5

acceptance/cmd/sandbox/register/success/test.toml

@@ -1,8 +1,9 @@
 # Acceptance tests for `databricks sandbox` subcommands.
 #
-# Ignore the local state file the CLI writes (sandbox.json) so it
-# doesn't bleed into the test directory's diff.
-Ignore = [".databricks"]
+# Ignore the local state file the CLI writes (sandbox.json) and the
+# generated SSH key + managed SSH config — register-path tests would
+# otherwise produce these on every run.
+Ignore = [".databricks", ".ssh"]
 
 # Sandbox commands are independent of the bundle engine. Override the
 # root EnvMatrix so the two engines don't both run in parallel for

acceptance/cmd/sandbox/test.toml

@@ -59,11 +59,14 @@ Examples:
 				return fmt.Errorf("failed to ensure sandbox SSH key: %w", err)
 			}
 
+			// Always print paths with forward slashes so acceptance
+			// goldens are stable across Linux/macOS/Windows.
+			displayKeyPath := filepath.ToSlash(keyPath)
 			stderr := cmd.ErrOrStderr()
 			if generated {
-				ok(ctx, "Generated SSH key at "+cmdio.Faint(ctx, keyPath))
+				ok(ctx, "Generated SSH key at "+cmdio.Faint(ctx, displayKeyPath))
 			} else {
-				field(ctx, stderr, "key", keyPath)
+				field(ctx, stderr, "key", displayKeyPath)
 			}
 
 			pubKeyData, err := os.ReadFile(keyPath + ".pub")
@@ -94,18 +97,26 @@ Examples:
 				_ = setGatewayHost(ctx, profile, registered.GatewayHost)
 			}
 
-			// Write the `Host <gateway>` block to ~/.ssh/config so
+			// Write the `Host <gateway>` block(s) to ~/.ssh/config so
 			// editor Remote-SSH ("Open in VS Code/Cursor") deep links
 			// and plain `ssh <id>@<gateway>` both work without the
-			// user pasting any config block. The gateway host comes
-			// straight from the registerKey response — no probe call
-			// needed, no heuristic.
-			if err := maybeWriteSSHConfig(ctx, keyPath, registered.GatewayHost); err != nil {
+			// user pasting any config block. Reads the full gateway
+			// set from local state so a user with profiles in
+			// multiple regions accumulates one block per gateway
+			// across registers, rather than overwriting.
+			if err := maybeWriteSSHConfig(ctx, keyPath); err != nil {
 				warn(ctx, fmt.Sprintf("Registered key, but failed to update ~/.ssh/config: %v", err))
 			}
 
+			// Tell the user what to do next based on whether they
+			// already have a sandbox: a brand-new register with no
+			// default yet should land on `create`, not `ssh`.
 			blank(stderr)
-			fmt.Fprintf(stderr, "  Run %s to connect.\n\n", cmdio.Bold(ctx, "databricks sandbox ssh"))
+			if getDefault(ctx, profile) == "" {
+				fmt.Fprintf(stderr, "  Run %s to provision your first sandbox.\n\n", cmdio.Bold(ctx, "databricks sandbox create"))
+			} else {
+				fmt.Fprintf(stderr, "  Run %s to connect.\n\n", cmdio.Bold(ctx, "databricks sandbox ssh"))
+			}
 			return nil
 		},
 	}

cmd/sandbox/register.go

@@ -17,24 +17,9 @@ import (
 	"github.com/spf13/cobra"
 )
 
-const (
-	defaultGatewayHost        = "uw2.dbrx.dev"
-	stagingDefaultGatewayHost = "ue1.s.dbrx.dev"
-	defaultGatewayPort        = "2222"
-)
-
-// resolveGatewayHost picks the SSH gateway hostname based on the workspace host.
-// Staging workspaces (*.staging.cloud.databricks.com etc.) route through
-// ue1.s.dbrx.dev; everything else uses uw2.dbrx.dev.
-func resolveGatewayHost(workspaceHost string) string {
-	if strings.Contains(workspaceHost, ".staging.") {
-		return stagingDefaultGatewayHost
-	}
-	return defaultGatewayHost
-}
+const defaultGatewayPort = "2222"
 
 func newSSHCommand() *cobra.Command {
-	var gatewayHost string
 	var gatewayPort string
 
 	cmd := &cobra.Command{
@@ -167,17 +152,17 @@ Examples:
 				}
 			}
 
-			// Resolution precedence: --gateway flag → fresh API response →
-			// cached value for this profile → workspace-host heuristic.
-			host := gatewayHost
-			if host == "" {
-				host = sandboxGatewayHost
-			}
+			// Resolution precedence: fresh API response → cached value
+			// for this profile. The server stamps the workspace's
+			// gateway on every Sandbox / SshKey response, so the cache
+			// is populated after the first call against the workspace
+			// (sandbox register / list / create / etc.).
+			host := sandboxGatewayHost
 			if host == "" {
 				host = getGatewayHost(ctx, profile)
 			}
 			if host == "" {
-				host = resolveGatewayHost(w.Config.Host)
+				return errors.New("could not resolve the sandbox SSH gateway — run `databricks sandbox register` first")
 			}
 
 			// Persist whatever the server just told us, so the next invocation
@@ -186,16 +171,17 @@ Examples:
 				_ = setGatewayHost(ctx, profile, sandboxGatewayHost)
 			}
 
-			// Don't print "Connected" here — ssh hasn't completed the
-			// handshake yet, so any success message would race ssh's
-			// own error output on the failure path.
-			s := spin(ctx, "Connecting to "+cmdio.Bold(ctx, sandboxID)+"…")
-			defer s.Close()
+			// No "Connecting…" spinner here. execSSHDirect replaces
+			// the CLI process via execv, so a deferred Close() would
+			// never run — and Bubble Tea hides the terminal cursor at
+			// init and only restores it on clean shutdown, leaving the
+			// cursor invisible for the entire ssh session and after it
+			// exits. The spinner couldn't animate during the handshake
+			// anyway (the CLI process is gone by then).
 			return execSSHDirect(sandboxID, host, gatewayPort, keyPath, extraArgs)
 		},
 	}
 
-	cmd.Flags().StringVar(&gatewayHost, "gateway", "", "Sandbox gateway hostname (auto-detected from profile if empty)")
 	cmd.Flags().StringVar(&gatewayPort, "port", defaultGatewayPort, "Sandbox gateway SSH port")
 
 	return cmd

cmd/sandbox/ssh.go

@@ -59,7 +59,7 @@ func sshConfigAlreadyManaged(ctx context.Context) (bool, error) {
 // writeSSHConfig writes the sandbox-managed SSH config block to a
 // managed file and, if not already present, adds an Include directive
 // to the user's ~/.ssh/config pointing at that file.
-func writeSSHConfig(ctx context.Context, keyPath, gatewayHost, gatewayPort string) (string, string, error) {
+func writeSSHConfig(ctx context.Context, keyPath string, gatewayHosts []string, gatewayPort string) (string, string, error) {
 	home, err := env.UserHomeDir(ctx)
 	if err != nil {
 		return "", "", err
@@ -72,7 +72,7 @@ func writeSSHConfig(ctx context.Context, keyPath, gatewayHost, gatewayPort strin
 	managedPath := filepath.Join(sshDir, sshIncludeFileName)
 	mainPath := filepath.Join(sshDir, "config")
 
-	block := buildSSHConfigBlock(keyPath, gatewayHost, gatewayPort)
+	block := buildSSHConfigBlock(keyPath, gatewayHosts, gatewayPort)
 	if err := writeManagedConfig(managedPath, block); err != nil {
 		return managedPath, mainPath, err
 	}
@@ -82,12 +82,12 @@ func writeSSHConfig(ctx context.Context, keyPath, gatewayHost, gatewayPort strin
 	return managedPath, mainPath, nil
 }
 
-// buildSSHConfigBlock renders the Host stanza we write to the
-// sandbox-managed include file. Mirrors the snippet the workspace UI's
-// "First time setup?" disclosure recommends — the Host key is the
-// literal gateway hostname (so editor Remote-SSH deep links like
-// `ssh-remote+<id>@<gateway>` resolve directly), and only the two
-// directives that meaningfully differ from SSH defaults are set:
+// buildSSHConfigBlock renders one Host stanza per gateway. Mirrors the
+// snippet the workspace UI's "First time setup?" disclosure recommends
+// — the Host key is the literal gateway hostname (so editor Remote-SSH
+// deep links like `ssh-remote+<id>@<gateway>` resolve directly), and
+// only the two directives that meaningfully differ from SSH defaults
+// are set:
 //
 //   - Port (the gateway listens on 2222, not 22)
 //   - IdentityFile + IdentitiesOnly (pin our key so ssh doesn't offer
@@ -99,13 +99,18 @@ func writeSSHConfig(ctx context.Context, keyPath, gatewayHost, gatewayPort strin
 // CLI's own `sandbox ssh` does via argv. No User directive either —
 // the per-sandbox identifier travels in the destination
 // (`ssh <sandbox-id>@<gateway>`).
-func buildSSHConfigBlock(keyPath, gatewayHost, gatewayPort string) string {
-	return fmt.Sprintf(`# Managed by `+"`databricks sandbox register`"+`. Manual edits will be overwritten.
-Host %s
-    Port %s
-    IdentityFile %s
-    IdentitiesOnly yes
-`, gatewayHost, gatewayPort, keyPath)
+//
+// One stanza per gateway means a user with profiles in multiple
+// regions doesn't lose IDE Remote-SSH for the earlier workspaces when
+// they `register` against a new one. Callers pass `gatewayHosts` from
+// state.allGatewayHosts.
+func buildSSHConfigBlock(keyPath string, gatewayHosts []string, gatewayPort string) string {
+	var b strings.Builder
+	b.WriteString("# Managed by `databricks sandbox register`. Manual edits will be overwritten.\n")
+	for _, gw := range gatewayHosts {
+		fmt.Fprintf(&b, "Host %s\n    Port %s\n    IdentityFile %s\n    IdentitiesOnly yes\n", gw, gatewayPort, keyPath)
+	}
+	return b.String()
 }
 
 // writeManagedConfig writes content to path atomically with 0600
@@ -183,14 +188,20 @@ func hasOurMarkedBlock(text string) bool {
 // maybeWriteSSHConfig writes the sandbox-managed SSH config, prompting
 // for consent the first time on this machine. Re-runs silently refresh
 // the managed file. Non-interactive contexts skip the write entirely.
-// The gateway host is the workspace-scoped one returned by the server
-// on the registerKey response — no heuristic, no probe call.
-func maybeWriteSSHConfig(ctx context.Context, keyPath, gatewayHost string) error {
-	if gatewayHost == "" {
-		// Server didn't stamp a gateway on the response. Old server
-		// without the gateway_host field, or genuinely no gateway
-		// available. Skip rather than guess.
-		cmdio.LogString(ctx, "  "+cmdio.Faint(ctx, "skipped SSH config update — server did not return a gateway host"))
+//
+// Reads the gateway-host set from state (populated by every API call
+// that touches a Sandbox or SshKey response, including the register
+// call the caller just made), so users with multiple workspaces in
+// different regions accumulate one Host stanza per unique gateway
+// instead of losing the earlier ones.
+func maybeWriteSSHConfig(ctx context.Context, keyPath string) error {
+	gateways := allGatewayHosts(ctx)
+	if len(gateways) == 0 {
+		// No gateway has been cached yet — register's API call must
+		// not have stamped one, or state is empty. Skip rather than
+		// guess; the next API call that returns a Sandbox/SshKey
+		// will populate the cache and a future register will write.
+		cmdio.LogString(ctx, "  "+cmdio.Faint(ctx, "skipped SSH config update — no sandbox gateway is known yet"))
 		return nil
 	}
 
@@ -204,12 +215,12 @@ func maybeWriteSSHConfig(ctx context.Context, keyPath, gatewayHost string) error
 			return err
 		}
 		if !cmdio.IsPromptSupported(ctx) {
-			cmdio.LogString(ctx, "  "+cmdio.Faint(ctx, "skipped SSH config update (non-interactive); re-run `databricks sandbox register` from a terminal to add the `Host "+gatewayHost+"` block"))
+			cmdio.LogString(ctx, "  "+cmdio.Faint(ctx, "skipped SSH config update (non-interactive); re-run `databricks sandbox register` from a terminal to add the sandbox gateway block(s)"))
 			return nil
 		}
 		question := fmt.Sprintf(
-			"Add a `Host %s` block to %s? This lets editor Remote-SSH (VS Code / Cursor) and `ssh <id>@%s` connect without further setup.",
-			gatewayHost, mainPath, gatewayHost,
+			"Add a sandbox `Host` block to %s? This lets editor Remote-SSH (VS Code / Cursor) and `ssh <id>@<gateway>` connect without further setup.",
+			mainPath,
 		)
 		confirmed, err := cmdio.AskYesOrNo(ctx, question)
 		if err != nil {
@@ -221,7 +232,7 @@ func maybeWriteSSHConfig(ctx context.Context, keyPath, gatewayHost string) error
 		}
 	}
 
-	managedPath, _, err := writeSSHConfig(ctx, keyPath, gatewayHost, defaultGatewayPort)
+	managedPath, _, err := writeSSHConfig(ctx, keyPath, gateways, defaultGatewayPort)
 	if err != nil {
 		return err
 	}

cmd/sandbox/sshconfig.go

@@ -1,18 +1,23 @@
 package sandbox
 
 import (
+	"io"
+	"io/fs"
 	"os"
 	"path/filepath"
 	"runtime"
 	"strings"
 	"testing"
 
+	"github.com/databricks/cli/libs/cmdio"
+	"github.com/databricks/cli/libs/env"
+	"github.com/databricks/cli/libs/flags"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
 func TestBuildSSHConfigBlockShape(t *testing.T) {
-	block := buildSSHConfigBlock("/home/u/.ssh/sandbox_ed25519", "uw2.dbrx.dev", "2222")
+	block := buildSSHConfigBlock("/home/u/.ssh/sandbox_ed25519", []string{"uw2.dbrx.dev"}, "2222")
 
 	// The Host key is the literal gateway hostname (so UI editor deep
 	// links like `ssh-remote+<id>@<gateway>` resolve directly), and we
@@ -47,6 +52,23 @@ func TestBuildSSHConfigBlockShape(t *testing.T) {
 	}
 }
 
+// Two registered workspaces in different regions must each get their
+// own Host stanza so editor Remote-SSH for one doesn't silently break
+// when the user registers against the other.
+func TestBuildSSHConfigBlockMultipleGateways(t *testing.T) {
+	block := buildSSHConfigBlock(
+		"/home/u/.ssh/sandbox_ed25519",
+		[]string{"uw2.dbrx.dev", "ue1.dbrx.dev"},
+		"2222",
+	)
+	assert.Contains(t, block, "Host uw2.dbrx.dev\n")
+	assert.Contains(t, block, "Host ue1.dbrx.dev\n")
+	// Each Host should be followed by Port + IdentityFile, so a naive
+	// "Port 2222" count is a quick sanity check on per-gateway repetition.
+	assert.Equal(t, 2, strings.Count(block, "Port 2222"))
+	assert.Equal(t, 2, strings.Count(block, "IdentityFile /home/u/.ssh/sandbox_ed25519"))
+}
+
 func TestWriteManagedConfigCreatesWithRightPerms(t *testing.T) {
 	dir := t.TempDir()
 	path := filepath.Join(dir, sshIncludeFileName)
@@ -143,6 +165,28 @@ func TestEnsureMainIncludesManagedPreservesUserConfigBelow(t *testing.T) {
 	assert.Less(t, beginIdx, userIdx, "managed block must precede the user's existing config")
 }
 
+// When state has no cached gateway hosts (no register has populated
+// the cache, or sandbox.json is missing), maybeWriteSSHConfig must
+// short-circuit before any disk I/O so it can't accidentally write a
+// malformed `Host \n` stanza. Also pins the no-disk-touch guarantee
+// against a future refactor that might `os.MkdirAll` before the check.
+func TestMaybeWriteSSHConfigSkipsWhenNoGatewaysCached(t *testing.T) {
+	home := t.TempDir()
+	ctx := env.WithUserHomeDir(t.Context(), home)
+	ctx = cmdio.InContext(ctx, cmdio.NewIO(ctx, flags.OutputText,
+		io.NopCloser(strings.NewReader("")), io.Discard, io.Discard, "", ""))
+
+	require.NoError(t, maybeWriteSSHConfig(ctx, filepath.Join(home, ".ssh", "sandbox_ed25519")))
+
+	// No managed file should have been created.
+	_, err := os.Stat(filepath.Join(home, ".ssh", sshIncludeFileName))
+	assert.ErrorIs(t, err, fs.ErrNotExist, "managed include file must not be created when no gateways are cached")
+
+	// And ~/.ssh/config must not have been touched either.
+	_, err = os.Stat(filepath.Join(home, ".ssh", "config"))
+	assert.ErrorIs(t, err, fs.ErrNotExist, "main ssh config must not be created when no gateways are cached")
+}
+
 func TestHasOurMarkedBlock(t *testing.T) {
 	cases := []struct {
 		name  string