Fix sandbox SSH during stop transition

cmd/sandbox/ssh.go

@@ -218,13 +218,27 @@ func verifyKeyRegistered(ctx context.Context, api *sandboxAPI, keyPath string) e
 	return fmt.Errorf("your sandbox SSH key (%s) is not registered with this workspace — run `databricks sandbox register` to re-register it", want)
 }
 
-// ensureRunning brings the named sandbox to Running with its own
-// spinner — caller must not have one open.
+// ensureRunning owns the lifecycle transition so SSH does not hang inside
+// connection setup. The caller must not have a spinner open.
 func ensureRunning(ctx context.Context, api *sandboxAPI, id, currentStatus string) (*sandboxEntry, error) {
+	if strings.EqualFold(currentStatus, "running") {
+		return &sandboxEntry{SandboxID: id, Status: currentStatus}, nil
+	}
+
 	s := spin(ctx, "Starting "+cmdio.Bold(ctx, id)+"…")
 	defer s.Close()
 
 	var sb *sandboxEntry
+	if strings.EqualFold(currentStatus, "stopping") {
+		stopped, err := waitForStopped(ctx, api, s, id)
+		if err != nil {
+			s.fail("Failed to start " + id)
+			return nil, err
+		}
+		currentStatus = stopped.Status
+		s.Update("Starting " + id + "…")
+	}
+
 	if strings.EqualFold(currentStatus, "stopped") {
 		updated, err := api.start(ctx, id)
 		if err != nil {

cmd/sandbox/ssh_test.go

@@ -1,11 +1,125 @@
 package sandbox
 
 import (
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"sync/atomic"
 	"testing"
 
+	"github.com/databricks/cli/libs/cmdio"
+	"github.com/databricks/databricks-sdk-go"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
+func TestEnsureRunningLifecycle(t *testing.T) {
+	type requestStep struct {
+		method string
+		path   string
+		status string
+	}
+
+	for _, tc := range []struct {
+		name          string
+		initialStatus string
+		steps         []requestStep
+		wantStatus    string
+		wantErr       string
+		wantStarts    int32
+	}{
+		{
+			name:          "stopping waits for stopped before starting",
+			initialStatus: "Stopping",
+			steps: []requestStep{
+				{method: http.MethodGet, path: sandboxPath("test-id"), status: "Stopped"},
+				{method: http.MethodPost, path: sandboxPath("test-id") + "/start", status: "Pending"},
+				{method: http.MethodGet, path: sandboxPath("test-id"), status: "Running"},
+			},
+			wantStatus: "Running",
+			wantStarts: 1,
+		},
+		{
+			name:          "stopped starts normally",
+			initialStatus: "Stopped",
+			steps: []requestStep{
+				{method: http.MethodPost, path: sandboxPath("test-id") + "/start", status: "Pending"},
+				{method: http.MethodGet, path: sandboxPath("test-id"), status: "Running"},
+			},
+			wantStatus: "Running",
+			wantStarts: 1,
+		},
+		{
+			name:          "running is a no-op",
+			initialStatus: "Running",
+			wantStatus:    "Running",
+		},
+		{
+			name:          "stopped after start remains an error",
+			initialStatus: "Stopped",
+			steps: []requestStep{
+				{method: http.MethodPost, path: sandboxPath("test-id") + "/start", status: "Pending"},
+				{method: http.MethodGet, path: sandboxPath("test-id"), status: "Stopped"},
+			},
+			wantErr:    `sandbox test-id reached unexpected state "Stopped" while starting`,
+			wantStarts: 1,
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			var requestCount atomic.Int32
+			var startCount atomic.Int32
+			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				if r.URL.Path == "/.well-known/databricks-config" {
+					assert.NoError(t, json.NewEncoder(w).Encode(map[string]string{
+						"oidc_endpoint": "https://workspace.example.test/oidc",
+					}))
+					return
+				}
+
+				index := int(requestCount.Add(1) - 1)
+				if index >= len(tc.steps) {
+					t.Errorf("unexpected request %s %s", r.Method, r.URL.Path)
+					http.Error(w, "unexpected request", http.StatusInternalServerError)
+					return
+				}
+
+				step := tc.steps[index]
+				assert.Equal(t, step.method, r.Method)
+				assert.Equal(t, step.path, r.URL.Path)
+				if r.Method == http.MethodPost && r.URL.Path == sandboxPath("test-id")+"/start" {
+					startCount.Add(1)
+				}
+				assert.NoError(t, json.NewEncoder(w).Encode(sandboxEntry{
+					SandboxID: "test-id",
+					Status:    step.status,
+				}))
+			}))
+			t.Cleanup(server.Close)
+
+			workspaceClient, err := databricks.NewWorkspaceClient(&databricks.Config{
+				Host:  server.URL,
+				Token: "test-token",
+			})
+			require.NoError(t, err)
+			api, err := newSandboxAPI(workspaceClient)
+			require.NoError(t, err)
+
+			ctx := cmdio.MockDiscard(t.Context())
+			got, err := ensureRunning(ctx, api, "test-id", tc.initialStatus)
+			if tc.wantErr != "" {
+				require.EqualError(t, err, tc.wantErr)
+				assert.Nil(t, got)
+			} else {
+				require.NoError(t, err)
+				require.NotNil(t, got)
+				assert.Equal(t, tc.wantStatus, got.Status)
+			}
+			assert.Equal(t, int32(len(tc.steps)), requestCount.Load())
+			assert.Equal(t, tc.wantStarts, startCount.Load())
+		})
+	}
+}
+
 func TestBuildSSHArgsBaseFlags(t *testing.T) {
 	got := buildSSHArgs("happy-panda-1234", "gw.example.test", "2222", "/keys/id", nil)
 	// Sanity: the destination is the last arg when no extras.

cmd/sandbox/start.go

@@ -93,29 +93,38 @@ Example:
 	return cmd
 }
 
-// waitForRunning polls the sandbox until it reaches Running or the timeout
-// elapses. The spinner is updated with the elapsed seconds each poll so the
-// user can tell progress is happening, and a transition to an unexpected
-// terminal state (Stopped / Terminated / Failed) short-circuits with a
-// useful error rather than waiting out the full timeout.
 func waitForRunning(ctx context.Context, api *sandboxAPI, s *spinner, id string) (*sandboxEntry, error) {
+	return waitForState(ctx, api, s, id, "Running", "Starting", "Stopped", "Terminated", "Failed")
+}
+
+// The API rejects start requests until teardown finishes, so STOPPING needs
+// its own polling phase.
+func waitForStopped(ctx context.Context, api *sandboxAPI, s *spinner, id string) (*sandboxEntry, error) {
+	return waitForState(ctx, api, s, id, "Stopped", "Stopping", "Terminated", "Failed")
+}
+
+// Centralizing lifecycle polling keeps timeout, cancellation, and terminal
+// state handling consistent across start and SSH.
+func waitForState(ctx context.Context, api *sandboxAPI, s *spinner, id, targetStatus, operation string, unexpectedStatuses ...string) (*sandboxEntry, error) {
 	start := time.Now()
 	deadline := start.Add(startWaitTimeout)
 	for {
 		sb, err := api.get(ctx, id)
 		if err != nil {
 			return nil, fmt.Errorf("polling status of %s: %w", id, err)
 		}
-		switch strings.ToLower(sb.Status) {
-		case "running":
+		if strings.EqualFold(sb.Status, targetStatus) {
 			return sb, nil
-		case "stopped", "terminated", "failed":
-			return nil, fmt.Errorf("sandbox %s reached unexpected state %q while starting", id, sb.Status)
+		}
+		for _, unexpectedStatus := range unexpectedStatuses {
+			if strings.EqualFold(sb.Status, unexpectedStatus) {
+				return nil, fmt.Errorf("sandbox %s reached unexpected state %q while %s", id, sb.Status, strings.ToLower(operation))
+			}
 		}
 		elapsed := time.Since(start).Round(time.Second)
-		s.Update(fmt.Sprintf("Starting %s… (%s)", id, elapsed))
+		s.Update(fmt.Sprintf("%s %s… (%s)", operation, id, elapsed))
 		if time.Now().After(deadline) {
-			return nil, fmt.Errorf("sandbox %s did not reach Running within %s (last seen %s)", id, startWaitTimeout, sb.Status)
+			return nil, fmt.Errorf("sandbox %s did not reach %s within %s (last seen %s)", id, targetStatus, startWaitTimeout, sb.Status)
 		}
 		select {
 		case <-ctx.Done():