From 756dfa2ebf729269206662d60e64370463bc1945 Mon Sep 17 00:00:00 2001
From: shuochen0311 <shuo.chen@databricks.com>
Date: Fri, 10 Apr 2026 18:20:26 +0000
Subject: [PATCH 01/36] Add Lakebox CLI for managing Databricks sandbox
 environments
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Lakebox provides SSH-accessible development environments backed by
microVM isolation. This adds CLI commands for lifecycle management:

- `lakebox auth login` — authenticate to a Databricks workspace
- `lakebox create` — create a new lakebox (with optional SSH public key)
- `lakebox list` — list your lakeboxes (shows status, key hash, default)
- `lakebox ssh` — SSH to your default lakebox (or create one on first use)
- `lakebox status <id>` — show lakebox details
- `lakebox delete <id>` — delete a lakebox
- `lakebox set-default <id>` — change the default lakebox

Features:
- Default lakebox management stored at ~/.databricks/lakebox.json per profile
- Automatic SSH config management (~/.ssh/config)
- Public key auth only (password/keyboard-interactive disabled in SSH config)
- Creates and sets default on first `lakebox ssh` if none exists
---
 cmd/cmd.go               | 126 +++++----------------
 cmd/lakebox/api.go       | 175 +++++++++++++++++++++++++++++
 cmd/lakebox/create.go    |  83 ++++++++++++++
 cmd/lakebox/default.go   |  39 +++++++
 cmd/lakebox/delete.go    |  51 +++++++++
 cmd/lakebox/exec_unix.go |  13 +++
 cmd/lakebox/lakebox.go   |  40 +++++++
 cmd/lakebox/list.go      |  70 ++++++++++++
 cmd/lakebox/ssh.go       | 235 +++++++++++++++++++++++++++++++++++++++
 cmd/lakebox/state.go     |  90 +++++++++++++++
 cmd/lakebox/status.go    |  58 ++++++++++
 11 files changed, 880 insertions(+), 100 deletions(-)
 create mode 100644 cmd/lakebox/api.go
 create mode 100644 cmd/lakebox/create.go
 create mode 100644 cmd/lakebox/default.go
 create mode 100644 cmd/lakebox/delete.go
 create mode 100644 cmd/lakebox/exec_unix.go
 create mode 100644 cmd/lakebox/lakebox.go
 create mode 100644 cmd/lakebox/list.go
 create mode 100644 cmd/lakebox/ssh.go
 create mode 100644 cmd/lakebox/state.go
 create mode 100644 cmd/lakebox/status.go

diff --git a/cmd/cmd.go b/cmd/cmd.go
index 014471f763..fe81149c08 100644
--- a/cmd/cmd.go
+++ b/cmd/cmd.go
@@ -2,117 +2,43 @@ package cmd
 
 import (
 	"context"
-	"strings"
 
-	"github.com/databricks/cli/cmd/psql"
-	ssh "github.com/databricks/cli/experimental/ssh/cmd"
-
-	"github.com/databricks/cli/cmd/account"
-	"github.com/databricks/cli/cmd/api"
 	"github.com/databricks/cli/cmd/auth"
-	"github.com/databricks/cli/cmd/bundle"
-	"github.com/databricks/cli/cmd/cache"
-	"github.com/databricks/cli/cmd/completion"
-	"github.com/databricks/cli/cmd/configure"
-	"github.com/databricks/cli/cmd/experimental"
-	"github.com/databricks/cli/cmd/fs"
-	"github.com/databricks/cli/cmd/labs"
-	"github.com/databricks/cli/cmd/pipelines"
+	"github.com/databricks/cli/cmd/lakebox"
 	"github.com/databricks/cli/cmd/root"
-	"github.com/databricks/cli/cmd/selftest"
-	"github.com/databricks/cli/cmd/sync"
-	"github.com/databricks/cli/cmd/version"
-	"github.com/databricks/cli/cmd/workspace"
-	"github.com/databricks/cli/libs/cmdgroup"
 	"github.com/spf13/cobra"
 )
 
-const (
-	mainGroup        = "main"
-	permissionsGroup = "permissions"
-)
-
-// configureGroups adds groups to the command, only if a group
-// has at least one available command.
-func configureGroups(cmd *cobra.Command, groups []cobra.Group) {
-	filteredGroups := cmdgroup.FilterGroups(groups, cmd.Commands())
-	for i := range filteredGroups {
-		cmd.AddGroup(&filteredGroups[i])
-	}
-}
-
-func accountCommand() *cobra.Command {
-	cmd := account.New()
-	configureGroups(cmd, account.Groups())
-	return cmd
-}
-
 func New(ctx context.Context) *cobra.Command {
 	cli := root.New(ctx)
+	cli.Use = "lakebox"
+	cli.Short = "Lakebox CLI — manage Databricks sandbox environments"
+	cli.Long = `Lakebox CLI — manage Databricks sandbox environments.
+
+Lakebox provides SSH-accessible development environments backed by
+microVM isolation. Each lakebox is a personal sandbox with pre-installed
+tooling (Python, Node.js, Rust, Databricks CLI) and persistent storage.
+
+Common workflows:
+  lakebox auth login                      # authenticate to Databricks
+  lakebox ssh                             # SSH to your default lakebox
+  lakebox ssh my-project                  # SSH to a named lakebox
+  lakebox list                            # list your lakeboxes
+  lakebox create                          # create a new lakebox
+  lakebox delete my-project               # delete a lakebox
+  lakebox status my-project               # show lakebox status
+
+The CLI manages your ~/.ssh/config so you can also connect directly:
+  ssh my-project                          # after 'lakebox ssh'
+`
+	cli.CompletionOptions.DisableDefaultCmd = true
 
-	// Add account subcommand.
-	cli.AddCommand(accountCommand())
-
-	// Add workspace subcommands.
-	workspaceCommands := workspace.All()
-	for _, cmd := range workspaceCommands {
-		// Order the permissions subcommands after the main commands.
-		for _, sub := range cmd.Commands() {
-			// some commands override groups in overrides.go, leave them as-is
-			if sub.GroupID != "" {
-				continue
-			}
-
-			switch {
-			case strings.HasSuffix(sub.Name(), "-permissions"), strings.HasSuffix(sub.Name(), "-permission-levels"):
-				sub.GroupID = permissionsGroup
-			default:
-				sub.GroupID = mainGroup
-			}
-		}
-
-		cli.AddCommand(cmd)
-
-		// Built-in groups for the workspace commands.
-		groups := []cobra.Group{
-			{
-				ID:    mainGroup,
-				Title: "Available Commands",
-			},
-			{
-				ID:    pipelines.ManagementGroupID,
-				Title: "Management Commands",
-			},
-			{
-				ID:    permissionsGroup,
-				Title: "Permission Commands",
-			},
-		}
-
-		configureGroups(cmd, groups)
-	}
-
-	// Add other subcommands.
-	cli.AddCommand(api.New())
 	cli.AddCommand(auth.New())
-	cli.AddCommand(completion.New())
-	cli.AddCommand(bundle.New())
-	cli.AddCommand(cache.New())
-	cli.AddCommand(experimental.New())
-	cli.AddCommand(psql.New())
-	cli.AddCommand(configure.New())
-	cli.AddCommand(fs.New())
-	cli.AddCommand(labs.New(ctx))
-	cli.AddCommand(sync.New())
-	cli.AddCommand(version.New())
-	cli.AddCommand(selftest.New())
-	cli.AddCommand(ssh.New())
 
-	// Add workspace command groups, filtering out empty groups or groups with only hidden commands.
-	configureGroups(cli, append(workspace.Groups(), cobra.Group{
-		ID:    "development",
-		Title: "Developer Tools",
-	}))
+	// Register lakebox subcommands directly at root level.
+	for _, sub := range lakebox.New().Commands() {
+		cli.AddCommand(sub)
+	}
 
 	return cli
 }
diff --git a/cmd/lakebox/api.go b/cmd/lakebox/api.go
new file mode 100644
index 0000000000..ff8f7d30b1
--- /dev/null
+++ b/cmd/lakebox/api.go
@@ -0,0 +1,175 @@
+package lakebox
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"strings"
+
+	"github.com/databricks/databricks-sdk-go"
+)
+
+const lakeboxAPIPath = "/api/2.0/lakebox"
+
+// lakeboxAPI wraps raw HTTP calls to the lakebox REST API.
+type lakeboxAPI struct {
+	w *databricks.WorkspaceClient
+}
+
+// createRequest is the JSON body for POST /api/2.0/lakebox.
+type createRequest struct {
+	PublicKey string `json:"public_key,omitempty"`
+}
+
+// createResponse is the JSON body returned by POST /api/2.0/lakebox.
+type createResponse struct {
+	LakeboxID string `json:"lakebox_id"`
+	Status    string `json:"status"`
+}
+
+// lakeboxEntry is a single item in the list response.
+type lakeboxEntry struct {
+	Name             string `json:"name"`
+	Status           string `json:"status"`
+	FQDN             string `json:"fqdn"`
+	PubkeyHashPrefix string `json:"pubkey_hash_prefix,omitempty"`
+}
+
+// listResponse is the JSON body returned by GET /api/2.0/lakebox.
+type listResponse struct {
+	Lakeboxes []lakeboxEntry `json:"lakeboxes"`
+}
+
+// apiError is the error body returned by the lakebox API.
+type apiError struct {
+	ErrorCode string `json:"error_code"`
+	Message   string `json:"message"`
+}
+
+func (e *apiError) Error() string {
+	return fmt.Sprintf("%s: %s", e.ErrorCode, e.Message)
+}
+
+func newLakeboxAPI(w *databricks.WorkspaceClient) *lakeboxAPI {
+	return &lakeboxAPI{w: w}
+}
+
+// create calls POST /api/2.0/lakebox with an optional public key.
+func (a *lakeboxAPI) create(ctx context.Context, publicKey string) (*createResponse, error) {
+	body := createRequest{PublicKey: publicKey}
+	jsonBody, err := json.Marshal(body)
+	if err != nil {
+		return nil, fmt.Errorf("failed to marshal request: %w", err)
+	}
+
+	resp, err := a.doRequest(ctx, "POST", lakeboxAPIPath, bytes.NewReader(jsonBody))
+	if err != nil {
+		return nil, err
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusCreated {
+		return nil, parseAPIError(resp)
+	}
+
+	var result createResponse
+	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
+		return nil, fmt.Errorf("failed to decode response: %w", err)
+	}
+	return &result, nil
+}
+
+// list calls GET /api/2.0/lakebox.
+func (a *lakeboxAPI) list(ctx context.Context) ([]lakeboxEntry, error) {
+	resp, err := a.doRequest(ctx, "GET", lakeboxAPIPath, nil)
+	if err != nil {
+		return nil, err
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return nil, parseAPIError(resp)
+	}
+
+	var result listResponse
+	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
+		return nil, fmt.Errorf("failed to decode response: %w", err)
+	}
+	return result.Lakeboxes, nil
+}
+
+// get calls GET /api/2.0/lakebox/{id}.
+func (a *lakeboxAPI) get(ctx context.Context, id string) (*lakeboxEntry, error) {
+	resp, err := a.doRequest(ctx, "GET", lakeboxAPIPath+"/"+id, nil)
+	if err != nil {
+		return nil, err
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return nil, parseAPIError(resp)
+	}
+
+	var result lakeboxEntry
+	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
+		return nil, fmt.Errorf("failed to decode response: %w", err)
+	}
+	return &result, nil
+}
+
+// delete calls DELETE /api/2.0/lakebox/{id}.
+func (a *lakeboxAPI) delete(ctx context.Context, id string) error {
+	resp, err := a.doRequest(ctx, "DELETE", lakeboxAPIPath+"/"+id, nil)
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusNoContent {
+		return parseAPIError(resp)
+	}
+	return nil
+}
+
+// doRequest makes an authenticated HTTP request to the workspace.
+func (a *lakeboxAPI) doRequest(ctx context.Context, method, path string, body io.Reader) (*http.Response, error) {
+	host := strings.TrimRight(a.w.Config.Host, "/")
+	url := host + path
+
+	req, err := http.NewRequestWithContext(ctx, method, url, body)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create request: %w", err)
+	}
+
+	if err := a.w.Config.Authenticate(req); err != nil {
+		return nil, fmt.Errorf("failed to authenticate: %w", err)
+	}
+
+	if body != nil {
+		req.Header.Set("Content-Type", "application/json")
+	}
+
+	return http.DefaultClient.Do(req)
+}
+
+func parseAPIError(resp *http.Response) error {
+	body, _ := io.ReadAll(resp.Body)
+	var apiErr apiError
+	if json.Unmarshal(body, &apiErr) == nil && apiErr.Message != "" {
+		return &apiErr
+	}
+	return fmt.Errorf("HTTP %d: %s", resp.StatusCode, string(body))
+}
+
+// extractLakeboxID extracts the short ID from a full resource name.
+// e.g. "apps/lakebox/instances/happy-panda-1234" -> "happy-panda-1234"
+func extractLakeboxID(name string) string {
+	parts := strings.Split(name, "/")
+	if len(parts) > 0 {
+		return parts[len(parts)-1]
+	}
+	return name
+}
diff --git a/cmd/lakebox/create.go b/cmd/lakebox/create.go
new file mode 100644
index 0000000000..872776cc8d
--- /dev/null
+++ b/cmd/lakebox/create.go
@@ -0,0 +1,83 @@
+package lakebox
+
+import (
+	"fmt"
+	"os"
+
+	"github.com/databricks/cli/cmd/root"
+	"github.com/databricks/cli/libs/cmdctx"
+	"github.com/spf13/cobra"
+)
+
+func newCreateCommand() *cobra.Command {
+	var publicKeyFile string
+
+	cmd := &cobra.Command{
+		Use:   "create",
+		Short: "Create a new Lakebox environment",
+		Long: `Create a new Lakebox environment.
+
+Creates a new personal development environment backed by a microVM.
+Blocks until the lakebox is running and prints the lakebox ID.
+
+If --public-key-file is provided, the key is installed in the lakebox's
+authorized_keys so you can SSH directly. Otherwise the gateway key is used.
+
+Example:
+  databricks lakebox create
+  databricks lakebox create --public-key-file ~/.ssh/id_ed25519.pub`,
+		PreRunE: root.MustWorkspaceClient,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			ctx := cmd.Context()
+			w := cmdctx.WorkspaceClient(ctx)
+			api := newLakeboxAPI(w)
+
+			var publicKey string
+			if publicKeyFile != "" {
+				data, err := os.ReadFile(publicKeyFile)
+				if err != nil {
+					return fmt.Errorf("failed to read public key file %s: %w", publicKeyFile, err)
+				}
+				publicKey = string(data)
+			}
+
+			fmt.Fprintf(cmd.ErrOrStderr(), "Creating lakebox...\n")
+
+			result, err := api.create(ctx, publicKey)
+			if err != nil {
+				return fmt.Errorf("failed to create lakebox: %w", err)
+			}
+
+			profile := w.Config.Profile
+			if profile == "" {
+				profile = w.Config.Host
+			}
+
+			// Set as default if no default exists, or the current default
+			// has been deleted (no longer in the list).
+			currentDefault := getDefault(profile)
+			shouldSetDefault := currentDefault == ""
+			if !shouldSetDefault && currentDefault != "" {
+				// Check if the current default still exists.
+				if _, err := api.get(ctx, currentDefault); err != nil {
+					shouldSetDefault = true
+				}
+			}
+			if shouldSetDefault {
+				if err := setDefault(profile, result.LakeboxID); err != nil {
+					fmt.Fprintf(cmd.ErrOrStderr(), "Warning: failed to save default: %v\n", err)
+				} else {
+					fmt.Fprintf(cmd.ErrOrStderr(), "Set as default lakebox.\n")
+				}
+			}
+
+			fmt.Fprintf(cmd.ErrOrStderr(), "Lakebox created (status: %s)\n", result.Status)
+			fmt.Fprintln(cmd.OutOrStdout(), result.LakeboxID)
+			return nil
+		},
+	}
+
+	cmd.Flags().StringVar(&publicKeyFile, "public-key-file", "", "Path to SSH public key file to install in the lakebox")
+
+	return cmd
+}
diff --git a/cmd/lakebox/default.go b/cmd/lakebox/default.go
new file mode 100644
index 0000000000..9d5a366c9c
--- /dev/null
+++ b/cmd/lakebox/default.go
@@ -0,0 +1,39 @@
+package lakebox
+
+import (
+	"fmt"
+
+	"github.com/databricks/cli/cmd/root"
+	"github.com/databricks/cli/libs/cmdctx"
+	"github.com/spf13/cobra"
+)
+
+func newSetDefaultCommand() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "set-default <lakebox-id>",
+		Short: "Set the default Lakebox for SSH",
+		Long: `Set the default Lakebox that 'databricks lakebox ssh' connects to.
+
+The default is stored locally in ~/.databricks/lakebox.json per profile.
+
+Example:
+  databricks lakebox set-default happy-panda-1234`,
+		Args:    cobra.ExactArgs(1),
+		PreRunE: root.MustWorkspaceClient,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			w := cmdctx.WorkspaceClient(cmd.Context())
+			profile := w.Config.Profile
+			if profile == "" {
+				profile = w.Config.Host
+			}
+
+			lakeboxID := args[0]
+			if err := setDefault(profile, lakeboxID); err != nil {
+				return fmt.Errorf("failed to set default: %w", err)
+			}
+			fmt.Fprintf(cmd.OutOrStdout(), "Default lakebox set to: %s\n", lakeboxID)
+			return nil
+		},
+	}
+	return cmd
+}
diff --git a/cmd/lakebox/delete.go b/cmd/lakebox/delete.go
new file mode 100644
index 0000000000..a814083ed3
--- /dev/null
+++ b/cmd/lakebox/delete.go
@@ -0,0 +1,51 @@
+package lakebox
+
+import (
+	"fmt"
+
+	"github.com/databricks/cli/cmd/root"
+	"github.com/databricks/cli/libs/cmdctx"
+	"github.com/spf13/cobra"
+)
+
+func newDeleteCommand() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "delete <lakebox-id>",
+		Short: "Delete a Lakebox environment",
+		Long: `Delete a Lakebox environment.
+
+Permanently terminates and removes the specified lakebox. Only the
+creator (same auth token) can delete a lakebox.
+
+Example:
+  databricks lakebox delete happy-panda-1234`,
+		Args:    cobra.ExactArgs(1),
+		PreRunE: root.MustWorkspaceClient,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			ctx := cmd.Context()
+			w := cmdctx.WorkspaceClient(ctx)
+			api := newLakeboxAPI(w)
+
+			lakeboxID := args[0]
+
+			if err := api.delete(ctx, lakeboxID); err != nil {
+				return fmt.Errorf("failed to delete lakebox %s: %w", lakeboxID, err)
+			}
+
+			// Clear default if we just deleted it.
+			profile := w.Config.Profile
+			if profile == "" {
+				profile = w.Config.Host
+			}
+			if getDefault(profile) == lakeboxID {
+				_ = clearDefault(profile)
+				fmt.Fprintf(cmd.ErrOrStderr(), "Cleared default lakebox.\n")
+			}
+
+			fmt.Fprintf(cmd.ErrOrStderr(), "Deleted lakebox %s\n", lakeboxID)
+			return nil
+		},
+	}
+
+	return cmd
+}
diff --git a/cmd/lakebox/exec_unix.go b/cmd/lakebox/exec_unix.go
new file mode 100644
index 0000000000..d47f629572
--- /dev/null
+++ b/cmd/lakebox/exec_unix.go
@@ -0,0 +1,13 @@
+//go:build !windows
+
+package lakebox
+
+import (
+	"os"
+	"syscall"
+)
+
+// execSyscall replaces the current process with the given command (Unix only).
+func execSyscall(path string, args []string) error {
+	return syscall.Exec(path, args, os.Environ())
+}
diff --git a/cmd/lakebox/lakebox.go b/cmd/lakebox/lakebox.go
new file mode 100644
index 0000000000..6523debef9
--- /dev/null
+++ b/cmd/lakebox/lakebox.go
@@ -0,0 +1,40 @@
+package lakebox
+
+import (
+	"github.com/spf13/cobra"
+)
+
+func New() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "lakebox",
+		Short: "Manage Databricks Lakebox environments",
+		Long: `Manage Databricks Lakebox environments.
+
+Lakebox provides SSH-accessible development environments backed by
+microVM isolation. Each lakebox is a personal sandbox with pre-installed
+tooling (Python, Node.js, Rust, Databricks CLI) and persistent storage.
+
+Common workflows:
+  databricks lakebox login                        # authenticate to Databricks
+  databricks lakebox ssh                          # SSH to your default lakebox
+  databricks lakebox ssh my-project               # SSH to a named lakebox
+  databricks lakebox list                         # list your lakeboxes
+  databricks lakebox create --name my-project     # create a new lakebox
+  databricks lakebox delete my-project            # delete a lakebox
+  databricks lakebox status                       # show current lakebox status
+
+The CLI manages your ~/.ssh/config so you can also connect directly:
+  ssh my-project                                  # after 'lakebox ssh --setup'
+`,
+	}
+
+	cmd.AddCommand(newLoginCommand())
+	cmd.AddCommand(newSSHCommand())
+	cmd.AddCommand(newListCommand())
+	cmd.AddCommand(newCreateCommand())
+	cmd.AddCommand(newDeleteCommand())
+	cmd.AddCommand(newStatusCommand())
+	cmd.AddCommand(newSetDefaultCommand())
+
+	return cmd
+}
diff --git a/cmd/lakebox/list.go b/cmd/lakebox/list.go
new file mode 100644
index 0000000000..bf80a9919e
--- /dev/null
+++ b/cmd/lakebox/list.go
@@ -0,0 +1,70 @@
+package lakebox
+
+import (
+	"encoding/json"
+	"fmt"
+
+	"github.com/databricks/cli/cmd/root"
+	"github.com/databricks/cli/libs/cmdctx"
+	"github.com/spf13/cobra"
+)
+
+func newListCommand() *cobra.Command {
+	var outputJSON bool
+
+	cmd := &cobra.Command{
+		Use:   "list",
+		Short: "List your Lakebox environments",
+		Long: `List your Lakebox environments.
+
+Shows all lakeboxes associated with your account, including their
+current status and ID.
+
+Example:
+  databricks lakebox list
+  databricks lakebox list --json`,
+		PreRunE: root.MustWorkspaceClient,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			ctx := cmd.Context()
+			w := cmdctx.WorkspaceClient(ctx)
+			api := newLakeboxAPI(w)
+
+			entries, err := api.list(ctx)
+			if err != nil {
+				return fmt.Errorf("failed to list lakeboxes: %w", err)
+			}
+
+			if outputJSON {
+				enc := json.NewEncoder(cmd.OutOrStdout())
+				enc.SetIndent("", "  ")
+				return enc.Encode(entries)
+			}
+
+			if len(entries) == 0 {
+				fmt.Fprintln(cmd.ErrOrStderr(), "No lakeboxes found.")
+				return nil
+			}
+
+			profile := w.Config.Profile
+			if profile == "" {
+				profile = w.Config.Host
+			}
+			defaultID := getDefault(profile)
+
+			fmt.Fprintf(cmd.OutOrStdout(), "  %-30s %-12s %-10s %s\n", "ID", "STATUS", "KEY", "DEFAULT")
+			for _, e := range entries {
+				id := extractLakeboxID(e.Name)
+				def := ""
+				if id == defaultID {
+					def = "*"
+				}
+				fmt.Fprintf(cmd.OutOrStdout(), "  %-30s %-12s %-10s %s\n", id, e.Status, e.PubkeyHashPrefix, def)
+			}
+			return nil
+		},
+	}
+
+	cmd.Flags().BoolVar(&outputJSON, "json", false, "Output as JSON")
+
+	return cmd
+}
diff --git a/cmd/lakebox/ssh.go b/cmd/lakebox/ssh.go
new file mode 100644
index 0000000000..1978dec684
--- /dev/null
+++ b/cmd/lakebox/ssh.go
@@ -0,0 +1,235 @@
+package lakebox
+
+import (
+	"fmt"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"runtime"
+	"strings"
+
+	"github.com/databricks/cli/cmd/root"
+	"github.com/databricks/cli/libs/cmdctx"
+	"github.com/spf13/cobra"
+)
+
+const (
+	defaultGatewayHost = "uw2.dbrx.dev"
+	defaultGatewayPort = "2222"
+
+	// SSH config block markers for idempotent updates.
+	sshConfigMarkerStart = "# --- Lakebox managed start ---"
+	sshConfigMarkerEnd   = "# --- Lakebox managed end ---"
+)
+
+func newSSHCommand() *cobra.Command {
+	var gatewayHost string
+	var gatewayPort string
+
+	cmd := &cobra.Command{
+		Use:   "ssh [lakebox-id]",
+		Short: "SSH into a Lakebox environment",
+		Long: `SSH into a Lakebox environment.
+
+This command:
+1. Authenticates to the Databricks workspace
+2. Ensures you have a local SSH key (~/.ssh/id_ed25519)
+3. Creates a lakebox if one doesn't exist (installs your public key)
+4. Updates ~/.ssh/config with a Host entry for the lakebox
+5. Connects via SSH using the lakebox ID as the SSH username
+
+Without arguments, creates a new lakebox. With a lakebox ID argument,
+connects to the specified lakebox.
+
+Example:
+  databricks lakebox ssh                    # create and connect to a new lakebox
+  databricks lakebox ssh happy-panda-1234   # connect to existing lakebox`,
+		Args: cobra.MaximumNArgs(1),
+		PreRunE: func(cmd *cobra.Command, args []string) error {
+			return root.MustWorkspaceClient(cmd, args)
+		},
+		RunE: func(cmd *cobra.Command, args []string) error {
+			ctx := cmd.Context()
+			w := cmdctx.WorkspaceClient(ctx)
+
+			profile := w.Config.Profile
+			if profile == "" {
+				profile = w.Config.Host
+			}
+
+			// Ensure SSH key exists.
+			keyPath, err := ensureSSHKey()
+			if err != nil {
+				return fmt.Errorf("failed to ensure SSH key: %w", err)
+			}
+			fmt.Fprintf(cmd.ErrOrStderr(), "Using SSH key: %s\n", keyPath)
+
+			// Determine lakebox ID:
+			// 1. Explicit arg → use it
+			// 2. Local default exists → use it
+			// 3. Neither → create a new one and set as default
+			var lakeboxID string
+			if len(args) > 0 {
+				lakeboxID = args[0]
+			} else if def := getDefault(profile); def != "" {
+				lakeboxID = def
+				fmt.Fprintf(cmd.ErrOrStderr(), "Using default lakebox: %s\n", lakeboxID)
+			} else {
+				api := newLakeboxAPI(w)
+				pubKeyData, err := os.ReadFile(keyPath + ".pub")
+				if err != nil {
+					return fmt.Errorf("failed to read public key %s.pub: %w", keyPath, err)
+				}
+
+				fmt.Fprintf(cmd.ErrOrStderr(), "Creating lakebox...\n")
+				result, err := api.create(ctx, string(pubKeyData))
+				if err != nil {
+					return fmt.Errorf("failed to create lakebox: %w", err)
+				}
+				lakeboxID = result.LakeboxID
+				fmt.Fprintf(cmd.ErrOrStderr(), "Lakebox %s created (status: %s)\n", lakeboxID, result.Status)
+
+				if err := setDefault(profile, lakeboxID); err != nil {
+					fmt.Fprintf(cmd.ErrOrStderr(), "Warning: failed to save default: %v\n", err)
+				}
+			}
+
+			// Write SSH config entry for this lakebox.
+			sshConfigPath, err := sshConfigFilePath()
+			if err != nil {
+				return err
+			}
+			entry := buildSSHConfigEntry(lakeboxID, gatewayHost, gatewayPort, keyPath)
+			if err := writeSSHConfigEntry(sshConfigPath, lakeboxID, entry); err != nil {
+				return fmt.Errorf("failed to update SSH config: %w", err)
+			}
+
+			fmt.Fprintf(cmd.ErrOrStderr(), "Connecting to %s@%s:%s...\n",
+				lakeboxID, gatewayHost, gatewayPort)
+			return execSSH(lakeboxID)
+		},
+	}
+
+	cmd.Flags().StringVar(&gatewayHost, "gateway", defaultGatewayHost, "Lakebox gateway hostname")
+	cmd.Flags().StringVar(&gatewayPort, "port", defaultGatewayPort, "Lakebox gateway SSH port")
+
+	return cmd
+}
+
+// ensureSSHKey checks for an existing SSH key and generates one if missing.
+func ensureSSHKey() (string, error) {
+	homeDir, err := os.UserHomeDir()
+	if err != nil {
+		return "", err
+	}
+
+	candidates := []string{
+		filepath.Join(homeDir, ".ssh", "id_ed25519"),
+		filepath.Join(homeDir, ".ssh", "id_rsa"),
+	}
+	for _, p := range candidates {
+		if _, err := os.Stat(p); err == nil {
+			return p, nil
+		}
+	}
+
+	// Generate ed25519 key.
+	keyPath := candidates[0]
+	sshDir := filepath.Dir(keyPath)
+	if err := os.MkdirAll(sshDir, 0700); err != nil {
+		return "", fmt.Errorf("failed to create %s: %w", sshDir, err)
+	}
+
+	cmd := exec.Command("ssh-keygen", "-t", "ed25519", "-f", keyPath, "-N", "", "-q")
+	cmd.Stdin = os.Stdin
+	cmd.Stdout = os.Stderr
+	cmd.Stderr = os.Stderr
+	if err := cmd.Run(); err != nil {
+		return "", fmt.Errorf("ssh-keygen failed: %w", err)
+	}
+
+	return keyPath, nil
+}
+
+func sshConfigFilePath() (string, error) {
+	homeDir, err := os.UserHomeDir()
+	if err != nil {
+		return "", err
+	}
+	return filepath.Join(homeDir, ".ssh", "config"), nil
+}
+
+// buildSSHConfigEntry creates the SSH config block for a lakebox.
+// The lakebox ID is used as both the Host alias and the SSH User.
+func buildSSHConfigEntry(lakeboxID, host, port, keyPath string) string {
+	return fmt.Sprintf(`Host %s
+    HostName %s
+    Port %s
+    User %s
+    IdentityFile %s
+    IdentitiesOnly yes
+    PreferredAuthentications publickey
+    PasswordAuthentication no
+    KbdInteractiveAuthentication no
+    StrictHostKeyChecking no
+    UserKnownHostsFile /dev/null
+    LogLevel INFO
+`, lakeboxID, host, port, lakeboxID, keyPath)
+}
+
+// writeSSHConfigEntry idempotently writes a single lakebox entry to ~/.ssh/config.
+// Replaces any existing lakebox block in-place.
+func writeSSHConfigEntry(configPath, lakeboxID, entry string) error {
+	sshDir := filepath.Dir(configPath)
+	if err := os.MkdirAll(sshDir, 0700); err != nil {
+		return err
+	}
+
+	existing, err := os.ReadFile(configPath)
+	if err != nil && !os.IsNotExist(err) {
+		return err
+	}
+
+	wrappedEntry := fmt.Sprintf("%s\n%s%s\n", sshConfigMarkerStart, entry, sshConfigMarkerEnd)
+	content := string(existing)
+
+	// Remove existing lakebox block if present.
+	startIdx := strings.Index(content, sshConfigMarkerStart)
+	if startIdx >= 0 {
+		endIdx := strings.Index(content[startIdx:], sshConfigMarkerEnd)
+		if endIdx >= 0 {
+			endIdx += startIdx + len(sshConfigMarkerEnd)
+			if endIdx < len(content) && content[endIdx] == '\n' {
+				endIdx++
+			}
+			content = content[:startIdx] + content[endIdx:]
+		}
+	}
+
+	if !strings.HasSuffix(content, "\n") && len(content) > 0 {
+		content += "\n"
+	}
+	content += wrappedEntry
+
+	return os.WriteFile(configPath, []byte(content), 0600)
+}
+
+// execSSH execs into ssh using the lakebox ID as the Host alias.
+func execSSH(lakeboxID string) error {
+	sshPath, err := exec.LookPath("ssh")
+	if err != nil {
+		return fmt.Errorf("ssh not found in PATH: %w", err)
+	}
+
+	args := []string{"ssh", lakeboxID}
+
+	if runtime.GOOS == "windows" {
+		cmd := exec.Command(sshPath, args[1:]...)
+		cmd.Stdin = os.Stdin
+		cmd.Stdout = os.Stdout
+		cmd.Stderr = os.Stderr
+		return cmd.Run()
+	}
+
+	return execSyscall(sshPath, args)
+}
diff --git a/cmd/lakebox/state.go b/cmd/lakebox/state.go
new file mode 100644
index 0000000000..c0c8ad2d84
--- /dev/null
+++ b/cmd/lakebox/state.go
@@ -0,0 +1,90 @@
+package lakebox
+
+import (
+	"encoding/json"
+	"fmt"
+	"os"
+	"path/filepath"
+)
+
+// stateFile stores per-profile lakebox defaults on the local filesystem.
+// Located at ~/.databricks/lakebox.json.
+type stateFile struct {
+	// Profile name → default lakebox ID.
+	Defaults map[string]string `json:"defaults"`
+}
+
+func stateFilePath() (string, error) {
+	home, err := os.UserHomeDir()
+	if err != nil {
+		return "", err
+	}
+	return filepath.Join(home, ".databricks", "lakebox.json"), nil
+}
+
+func loadState() (*stateFile, error) {
+	path, err := stateFilePath()
+	if err != nil {
+		return nil, err
+	}
+
+	data, err := os.ReadFile(path)
+	if os.IsNotExist(err) {
+		return &stateFile{Defaults: make(map[string]string)}, nil
+	}
+	if err != nil {
+		return nil, fmt.Errorf("failed to read %s: %w", path, err)
+	}
+
+	var state stateFile
+	if err := json.Unmarshal(data, &state); err != nil {
+		return &stateFile{Defaults: make(map[string]string)}, nil
+	}
+	if state.Defaults == nil {
+		state.Defaults = make(map[string]string)
+	}
+	return &state, nil
+}
+
+func saveState(state *stateFile) error {
+	path, err := stateFilePath()
+	if err != nil {
+		return err
+	}
+
+	if err := os.MkdirAll(filepath.Dir(path), 0700); err != nil {
+		return err
+	}
+
+	data, err := json.MarshalIndent(state, "", "  ")
+	if err != nil {
+		return err
+	}
+	return os.WriteFile(path, data, 0600)
+}
+
+func getDefault(profile string) string {
+	state, err := loadState()
+	if err != nil {
+		return ""
+	}
+	return state.Defaults[profile]
+}
+
+func setDefault(profile, lakeboxID string) error {
+	state, err := loadState()
+	if err != nil {
+		return err
+	}
+	state.Defaults[profile] = lakeboxID
+	return saveState(state)
+}
+
+func clearDefault(profile string) error {
+	state, err := loadState()
+	if err != nil {
+		return err
+	}
+	delete(state.Defaults, profile)
+	return saveState(state)
+}
diff --git a/cmd/lakebox/status.go b/cmd/lakebox/status.go
new file mode 100644
index 0000000000..1afd968211
--- /dev/null
+++ b/cmd/lakebox/status.go
@@ -0,0 +1,58 @@
+package lakebox
+
+import (
+	"encoding/json"
+	"fmt"
+
+	"github.com/databricks/cli/cmd/root"
+	"github.com/databricks/cli/libs/cmdctx"
+	"github.com/spf13/cobra"
+)
+
+func newStatusCommand() *cobra.Command {
+	var outputJSON bool
+
+	cmd := &cobra.Command{
+		Use:   "status <lakebox-id>",
+		Short: "Show Lakebox environment status",
+		Long: `Show detailed status of a Lakebox environment.
+
+Example:
+  databricks lakebox status happy-panda-1234
+  databricks lakebox status happy-panda-1234 --json`,
+		Args:    cobra.ExactArgs(1),
+		PreRunE: root.MustWorkspaceClient,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			ctx := cmd.Context()
+			w := cmdctx.WorkspaceClient(ctx)
+			api := newLakeboxAPI(w)
+
+			lakeboxID := args[0]
+
+			entry, err := api.get(ctx, lakeboxID)
+			if err != nil {
+				return fmt.Errorf("failed to get lakebox %s: %w", lakeboxID, err)
+			}
+
+			if outputJSON {
+				enc := json.NewEncoder(cmd.OutOrStdout())
+				enc.SetIndent("", "  ")
+				return enc.Encode(entry)
+			}
+
+			fmt.Fprintf(cmd.OutOrStdout(), "ID:     %s\n", extractLakeboxID(entry.Name))
+			fmt.Fprintf(cmd.OutOrStdout(), "Status: %s\n", entry.Status)
+			if entry.FQDN != "" {
+				fmt.Fprintf(cmd.OutOrStdout(), "FQDN:   %s\n", entry.FQDN)
+			}
+			if entry.PubkeyHashPrefix != "" {
+				fmt.Fprintf(cmd.OutOrStdout(), "Key:    %s\n", entry.PubkeyHashPrefix)
+			}
+			return nil
+		},
+	}
+
+	cmd.Flags().BoolVar(&outputJSON, "json", false, "Output as JSON")
+
+	return cmd
+}

From c20c6dfaa65e5db081292d051fd5f45517ff6c1a Mon Sep 17 00:00:00 2001
From: shuochen0311 <shuo.chen@databricks.com>
Date: Mon, 13 Apr 2026 20:29:55 +0000
Subject: [PATCH 02/36] Remove KEY column from list, add register-key command

- Remove PubkeyHashPrefix field from lakeboxEntry (no longer returned by API)
- Remove KEY column from list output
- Remove Key line from status output
- Add register-key subcommand for SSH public key registration

Co-authored-by: Isaac
---
 cmd/lakebox/api.go          | 32 ++++++++++++++++++---
 cmd/lakebox/lakebox.go      | 19 +++++++------
 cmd/lakebox/list.go         |  4 +--
 cmd/lakebox/register_key.go | 55 +++++++++++++++++++++++++++++++++++++
 cmd/lakebox/status.go       |  3 --
 5 files changed, 95 insertions(+), 18 deletions(-)
 create mode 100644 cmd/lakebox/register_key.go

diff --git a/cmd/lakebox/api.go b/cmd/lakebox/api.go
index ff8f7d30b1..94877b4a42 100644
--- a/cmd/lakebox/api.go
+++ b/cmd/lakebox/api.go
@@ -32,10 +32,9 @@ type createResponse struct {
 
 // lakeboxEntry is a single item in the list response.
 type lakeboxEntry struct {
-	Name             string `json:"name"`
-	Status           string `json:"status"`
-	FQDN             string `json:"fqdn"`
-	PubkeyHashPrefix string `json:"pubkey_hash_prefix,omitempty"`
+	Name   string `json:"name"`
+	Status string `json:"status"`
+	FQDN   string `json:"fqdn"`
 }
 
 // listResponse is the JSON body returned by GET /api/2.0/lakebox.
@@ -164,6 +163,31 @@ func parseAPIError(resp *http.Response) error {
 	return fmt.Errorf("HTTP %d: %s", resp.StatusCode, string(body))
 }
 
+// registerKeyRequest is the JSON body for POST /api/2.0/lakebox/register-key.
+type registerKeyRequest struct {
+	PublicKey string `json:"public_key"`
+}
+
+// registerKey calls POST /api/2.0/lakebox/register-key.
+func (a *lakeboxAPI) registerKey(ctx context.Context, publicKey string) error {
+	body := registerKeyRequest{PublicKey: publicKey}
+	jsonBody, err := json.Marshal(body)
+	if err != nil {
+		return fmt.Errorf("failed to marshal request: %w", err)
+	}
+
+	resp, err := a.doRequest(ctx, "POST", lakeboxAPIPath+"/register-key", bytes.NewReader(jsonBody))
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return parseAPIError(resp)
+	}
+	return nil
+}
+
 // extractLakeboxID extracts the short ID from a full resource name.
 // e.g. "apps/lakebox/instances/happy-panda-1234" -> "happy-panda-1234"
 func extractLakeboxID(name string) string {
diff --git a/cmd/lakebox/lakebox.go b/cmd/lakebox/lakebox.go
index 6523debef9..aa9463bca8 100644
--- a/cmd/lakebox/lakebox.go
+++ b/cmd/lakebox/lakebox.go
@@ -15,26 +15,27 @@ microVM isolation. Each lakebox is a personal sandbox with pre-installed
 tooling (Python, Node.js, Rust, Databricks CLI) and persistent storage.
 
 Common workflows:
-  databricks lakebox login                        # authenticate to Databricks
-  databricks lakebox ssh                          # SSH to your default lakebox
-  databricks lakebox ssh my-project               # SSH to a named lakebox
-  databricks lakebox list                         # list your lakeboxes
-  databricks lakebox create --name my-project     # create a new lakebox
-  databricks lakebox delete my-project            # delete a lakebox
-  databricks lakebox status                       # show current lakebox status
+  lakebox auth login                              # authenticate to Databricks
+  lakebox ssh                                     # SSH to your default lakebox
+  lakebox ssh my-project                          # SSH to a named lakebox
+  lakebox list                                    # list your lakeboxes
+  lakebox create                                  # create a new lakebox
+  lakebox delete my-project                       # delete a lakebox
+  lakebox status my-project                       # show lakebox status
+  lakebox register-key --public-key-file ~/.ssh/id_rsa.pub  # register SSH key
 
 The CLI manages your ~/.ssh/config so you can also connect directly:
-  ssh my-project                                  # after 'lakebox ssh --setup'
+  ssh my-project                                  # after 'lakebox ssh'
 `,
 	}
 
-	cmd.AddCommand(newLoginCommand())
 	cmd.AddCommand(newSSHCommand())
 	cmd.AddCommand(newListCommand())
 	cmd.AddCommand(newCreateCommand())
 	cmd.AddCommand(newDeleteCommand())
 	cmd.AddCommand(newStatusCommand())
 	cmd.AddCommand(newSetDefaultCommand())
+	cmd.AddCommand(newRegisterKeyCommand())
 
 	return cmd
 }
diff --git a/cmd/lakebox/list.go b/cmd/lakebox/list.go
index bf80a9919e..90139d6be8 100644
--- a/cmd/lakebox/list.go
+++ b/cmd/lakebox/list.go
@@ -51,14 +51,14 @@ Example:
 			}
 			defaultID := getDefault(profile)
 
-			fmt.Fprintf(cmd.OutOrStdout(), "  %-30s %-12s %-10s %s\n", "ID", "STATUS", "KEY", "DEFAULT")
+			fmt.Fprintf(cmd.OutOrStdout(), "  %-30s %-12s %s\n", "ID", "STATUS", "DEFAULT")
 			for _, e := range entries {
 				id := extractLakeboxID(e.Name)
 				def := ""
 				if id == defaultID {
 					def = "*"
 				}
-				fmt.Fprintf(cmd.OutOrStdout(), "  %-30s %-12s %-10s %s\n", id, e.Status, e.PubkeyHashPrefix, def)
+				fmt.Fprintf(cmd.OutOrStdout(), "  %-30s %-12s %s\n", id, e.Status, def)
 			}
 			return nil
 		},
diff --git a/cmd/lakebox/register_key.go b/cmd/lakebox/register_key.go
new file mode 100644
index 0000000000..5a19cc4f57
--- /dev/null
+++ b/cmd/lakebox/register_key.go
@@ -0,0 +1,55 @@
+package lakebox
+
+import (
+	"fmt"
+	"os"
+
+	"github.com/databricks/cli/cmd/root"
+	"github.com/databricks/cli/libs/cmdctx"
+	"github.com/spf13/cobra"
+)
+
+func newRegisterKeyCommand() *cobra.Command {
+	var publicKeyFile string
+
+	cmd := &cobra.Command{
+		Use:   "register-key",
+		Short: "Register an SSH public key for lakebox access",
+		Long: `Register an SSH public key with the lakebox service.
+
+Once registered, the key can be used to SSH into any of your lakeboxes.
+A user can have multiple registered keys; any of them grants access to
+all lakeboxes owned by that user.
+
+Example:
+  databricks lakebox register-key --public-key-file ~/.ssh/id_ed25519.pub`,
+		PreRunE: root.MustWorkspaceClient,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			ctx := cmd.Context()
+			w := cmdctx.WorkspaceClient(ctx)
+			api := newLakeboxAPI(w)
+
+			if publicKeyFile == "" {
+				return fmt.Errorf("--public-key-file is required")
+			}
+
+			data, err := os.ReadFile(publicKeyFile)
+			if err != nil {
+				return fmt.Errorf("failed to read public key file %s: %w", publicKeyFile, err)
+			}
+
+			publicKey := string(data)
+			if err := api.registerKey(ctx, publicKey); err != nil {
+				return fmt.Errorf("failed to register key: %w", err)
+			}
+
+			fmt.Fprintln(cmd.ErrOrStderr(), "SSH public key registered.")
+			return nil
+		},
+	}
+
+	cmd.Flags().StringVar(&publicKeyFile, "public-key-file", "", "Path to SSH public key file to register")
+	_ = cmd.MarkFlagRequired("public-key-file")
+
+	return cmd
+}
diff --git a/cmd/lakebox/status.go b/cmd/lakebox/status.go
index 1afd968211..4bb130496d 100644
--- a/cmd/lakebox/status.go
+++ b/cmd/lakebox/status.go
@@ -45,9 +45,6 @@ Example:
 			if entry.FQDN != "" {
 				fmt.Fprintf(cmd.OutOrStdout(), "FQDN:   %s\n", entry.FQDN)
 			}
-			if entry.PubkeyHashPrefix != "" {
-				fmt.Fprintf(cmd.OutOrStdout(), "Key:    %s\n", entry.PubkeyHashPrefix)
-			}
 			return nil
 		},
 	}

From f8f8cc1aa04add672a448c1b399589ecb1a49435 Mon Sep 17 00:00:00 2001
From: shuochen0311 <shuo.chen@databricks.com>
Date: Tue, 14 Apr 2026 03:11:22 +0000
Subject: [PATCH 03/36] Simplify SSH flow: register command, direct SSH args,
 remove config writes

- Add 'register' command: generates ~/.ssh/lakebox_rsa and registers with API
- Remove 'register-key' command (replaced by 'register')
- Remove 'login' command (use 'auth login' + 'register' separately)
- SSH command passes options directly as args instead of writing ~/.ssh/config
- Check for ssh-keygen availability with helpful install instructions

Co-authored-by: Isaac
---
 cmd/cmd.go                  |   6 +-
 cmd/lakebox/lakebox.go      |  23 +++---
 cmd/lakebox/register.go     | 110 ++++++++++++++++++++++++++++
 cmd/lakebox/register_key.go |  55 --------------
 cmd/lakebox/ssh.go          | 141 +++++-------------------------------
 5 files changed, 148 insertions(+), 187 deletions(-)
 create mode 100644 cmd/lakebox/register.go
 delete mode 100644 cmd/lakebox/register_key.go

diff --git a/cmd/cmd.go b/cmd/cmd.go
index fe81149c08..c120f25aa7 100644
--- a/cmd/cmd.go
+++ b/cmd/cmd.go
@@ -19,8 +19,12 @@ Lakebox provides SSH-accessible development environments backed by
 microVM isolation. Each lakebox is a personal sandbox with pre-installed
 tooling (Python, Node.js, Rust, Databricks CLI) and persistent storage.
 
+Getting started:
+  lakebox auth login --host https://...   # authenticate to Databricks
+  lakebox register                        # generate SSH key and register
+  lakebox ssh                             # SSH to your default lakebox
+
 Common workflows:
-  lakebox auth login                      # authenticate to Databricks
   lakebox ssh                             # SSH to your default lakebox
   lakebox ssh my-project                  # SSH to a named lakebox
   lakebox list                            # list your lakeboxes
diff --git a/cmd/lakebox/lakebox.go b/cmd/lakebox/lakebox.go
index aa9463bca8..127b5d93bf 100644
--- a/cmd/lakebox/lakebox.go
+++ b/cmd/lakebox/lakebox.go
@@ -14,28 +14,31 @@ Lakebox provides SSH-accessible development environments backed by
 microVM isolation. Each lakebox is a personal sandbox with pre-installed
 tooling (Python, Node.js, Rust, Databricks CLI) and persistent storage.
 
+Getting started:
+  lakebox auth login --host https://...   # authenticate to Databricks
+  lakebox register                        # generate SSH key and register
+  lakebox ssh                             # SSH to your default lakebox
+
 Common workflows:
-  lakebox auth login                              # authenticate to Databricks
-  lakebox ssh                                     # SSH to your default lakebox
-  lakebox ssh my-project                          # SSH to a named lakebox
-  lakebox list                                    # list your lakeboxes
-  lakebox create                                  # create a new lakebox
-  lakebox delete my-project                       # delete a lakebox
-  lakebox status my-project                       # show lakebox status
-  lakebox register-key --public-key-file ~/.ssh/id_rsa.pub  # register SSH key
+  lakebox ssh                             # SSH to your default lakebox
+  lakebox ssh my-project                  # SSH to a named lakebox
+  lakebox list                            # list your lakeboxes
+  lakebox create                          # create a new lakebox
+  lakebox delete my-project               # delete a lakebox
+  lakebox status my-project               # show lakebox status
 
 The CLI manages your ~/.ssh/config so you can also connect directly:
-  ssh my-project                                  # after 'lakebox ssh'
+  ssh my-project                          # after 'lakebox ssh'
 `,
 	}
 
+	cmd.AddCommand(newRegisterCommand())
 	cmd.AddCommand(newSSHCommand())
 	cmd.AddCommand(newListCommand())
 	cmd.AddCommand(newCreateCommand())
 	cmd.AddCommand(newDeleteCommand())
 	cmd.AddCommand(newStatusCommand())
 	cmd.AddCommand(newSetDefaultCommand())
-	cmd.AddCommand(newRegisterKeyCommand())
 
 	return cmd
 }
diff --git a/cmd/lakebox/register.go b/cmd/lakebox/register.go
new file mode 100644
index 0000000000..7286a14bf5
--- /dev/null
+++ b/cmd/lakebox/register.go
@@ -0,0 +1,110 @@
+package lakebox
+
+import (
+	"fmt"
+	"os"
+	"os/exec"
+	"path/filepath"
+
+	"github.com/databricks/cli/cmd/root"
+	"github.com/databricks/cli/libs/cmdctx"
+	"github.com/spf13/cobra"
+)
+
+const lakeboxKeyName = "lakebox_rsa"
+
+func newRegisterCommand() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "register",
+		Short: "Register this machine for lakebox SSH access",
+		Long: `Generate a dedicated SSH key for lakebox and register it with the service.
+
+This command:
+1. Generates an RSA SSH key at ~/.ssh/lakebox_rsa (if it doesn't exist)
+2. Registers the public key with the lakebox service
+
+After registration, 'lakebox ssh' will use this key automatically.
+Run this once per machine.
+
+Example:
+  lakebox register`,
+		PreRunE: root.MustWorkspaceClient,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			ctx := cmd.Context()
+			w := cmdctx.WorkspaceClient(ctx)
+			api := newLakeboxAPI(w)
+
+			keyPath, generated, err := ensureLakeboxKey()
+			if err != nil {
+				return fmt.Errorf("failed to ensure lakebox SSH key: %w", err)
+			}
+
+			if generated {
+				fmt.Fprintf(cmd.ErrOrStderr(), "Generated SSH key: %s\n", keyPath)
+			} else {
+				fmt.Fprintf(cmd.ErrOrStderr(), "Using existing SSH key: %s\n", keyPath)
+			}
+
+			pubKeyData, err := os.ReadFile(keyPath + ".pub")
+			if err != nil {
+				return fmt.Errorf("failed to read public key %s.pub: %w", keyPath, err)
+			}
+
+			if err := api.registerKey(ctx, string(pubKeyData)); err != nil {
+				return fmt.Errorf("failed to register key: %w", err)
+			}
+
+			fmt.Fprintln(cmd.ErrOrStderr(), "Registered. You can now use 'lakebox ssh' to connect.")
+			return nil
+		},
+	}
+
+	return cmd
+}
+
+// lakeboxKeyPath returns the path to the dedicated lakebox SSH key.
+func lakeboxKeyPath() (string, error) {
+	homeDir, err := os.UserHomeDir()
+	if err != nil {
+		return "", err
+	}
+	return filepath.Join(homeDir, ".ssh", lakeboxKeyName), nil
+}
+
+// ensureLakeboxKey returns the path to the lakebox SSH key, generating it if
+// it doesn't exist. Returns (path, wasGenerated, error).
+func ensureLakeboxKey() (string, bool, error) {
+	keyPath, err := lakeboxKeyPath()
+	if err != nil {
+		return "", false, err
+	}
+
+	if _, err := os.Stat(keyPath); err == nil {
+		return keyPath, false, nil
+	}
+
+	// Check that ssh-keygen is available before trying to generate.
+	if _, err := exec.LookPath("ssh-keygen"); err != nil {
+		return "", false, fmt.Errorf(
+			"ssh-keygen not found in PATH.\n" +
+				"Please install OpenSSH and run 'lakebox register' again.\n" +
+				"  macOS:   brew install openssh\n" +
+				"  Ubuntu:  sudo apt install openssh-client\n" +
+				"  Windows: install Git for Windows (includes ssh-keygen)")
+	}
+
+	sshDir := filepath.Dir(keyPath)
+	if err := os.MkdirAll(sshDir, 0700); err != nil {
+		return "", false, fmt.Errorf("failed to create %s: %w", sshDir, err)
+	}
+
+	genCmd := exec.Command("ssh-keygen", "-t", "rsa", "-b", "4096", "-f", keyPath, "-N", "", "-q", "-C", "lakebox")
+	genCmd.Stdin = os.Stdin
+	genCmd.Stdout = os.Stderr
+	genCmd.Stderr = os.Stderr
+	if err := genCmd.Run(); err != nil {
+		return "", false, fmt.Errorf("ssh-keygen failed: %w", err)
+	}
+
+	return keyPath, true, nil
+}
diff --git a/cmd/lakebox/register_key.go b/cmd/lakebox/register_key.go
deleted file mode 100644
index 5a19cc4f57..0000000000
--- a/cmd/lakebox/register_key.go
+++ /dev/null
@@ -1,55 +0,0 @@
-package lakebox
-
-import (
-	"fmt"
-	"os"
-
-	"github.com/databricks/cli/cmd/root"
-	"github.com/databricks/cli/libs/cmdctx"
-	"github.com/spf13/cobra"
-)
-
-func newRegisterKeyCommand() *cobra.Command {
-	var publicKeyFile string
-
-	cmd := &cobra.Command{
-		Use:   "register-key",
-		Short: "Register an SSH public key for lakebox access",
-		Long: `Register an SSH public key with the lakebox service.
-
-Once registered, the key can be used to SSH into any of your lakeboxes.
-A user can have multiple registered keys; any of them grants access to
-all lakeboxes owned by that user.
-
-Example:
-  databricks lakebox register-key --public-key-file ~/.ssh/id_ed25519.pub`,
-		PreRunE: root.MustWorkspaceClient,
-		RunE: func(cmd *cobra.Command, args []string) error {
-			ctx := cmd.Context()
-			w := cmdctx.WorkspaceClient(ctx)
-			api := newLakeboxAPI(w)
-
-			if publicKeyFile == "" {
-				return fmt.Errorf("--public-key-file is required")
-			}
-
-			data, err := os.ReadFile(publicKeyFile)
-			if err != nil {
-				return fmt.Errorf("failed to read public key file %s: %w", publicKeyFile, err)
-			}
-
-			publicKey := string(data)
-			if err := api.registerKey(ctx, publicKey); err != nil {
-				return fmt.Errorf("failed to register key: %w", err)
-			}
-
-			fmt.Fprintln(cmd.ErrOrStderr(), "SSH public key registered.")
-			return nil
-		},
-	}
-
-	cmd.Flags().StringVar(&publicKeyFile, "public-key-file", "", "Path to SSH public key file to register")
-	_ = cmd.MarkFlagRequired("public-key-file")
-
-	return cmd
-}
diff --git a/cmd/lakebox/ssh.go b/cmd/lakebox/ssh.go
index 1978dec684..8868f38e81 100644
--- a/cmd/lakebox/ssh.go
+++ b/cmd/lakebox/ssh.go
@@ -4,9 +4,7 @@ import (
 	"fmt"
 	"os"
 	"os/exec"
-	"path/filepath"
 	"runtime"
-	"strings"
 
 	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdctx"
@@ -16,10 +14,6 @@ import (
 const (
 	defaultGatewayHost = "uw2.dbrx.dev"
 	defaultGatewayPort = "2222"
-
-	// SSH config block markers for idempotent updates.
-	sshConfigMarkerStart = "# --- Lakebox managed start ---"
-	sshConfigMarkerEnd   = "# --- Lakebox managed end ---"
 )
 
 func newSSHCommand() *cobra.Command {
@@ -57,10 +51,13 @@ Example:
 				profile = w.Config.Host
 			}
 
-			// Ensure SSH key exists.
-			keyPath, err := ensureSSHKey()
+			// Use the dedicated lakebox SSH key.
+			keyPath, err := lakeboxKeyPath()
 			if err != nil {
-				return fmt.Errorf("failed to ensure SSH key: %w", err)
+				return fmt.Errorf("failed to determine lakebox key path: %w", err)
+			}
+			if _, err := os.Stat(keyPath); os.IsNotExist(err) {
+				return fmt.Errorf("lakebox SSH key not found at %s — run 'lakebox register' first", keyPath)
 			}
 			fmt.Fprintf(cmd.ErrOrStderr(), "Using SSH key: %s\n", keyPath)
 
@@ -94,19 +91,9 @@ Example:
 				}
 			}
 
-			// Write SSH config entry for this lakebox.
-			sshConfigPath, err := sshConfigFilePath()
-			if err != nil {
-				return err
-			}
-			entry := buildSSHConfigEntry(lakeboxID, gatewayHost, gatewayPort, keyPath)
-			if err := writeSSHConfigEntry(sshConfigPath, lakeboxID, entry); err != nil {
-				return fmt.Errorf("failed to update SSH config: %w", err)
-			}
-
 			fmt.Fprintf(cmd.ErrOrStderr(), "Connecting to %s@%s:%s...\n",
 				lakeboxID, gatewayHost, gatewayPort)
-			return execSSH(lakeboxID)
+			return execSSHDirect(lakeboxID, gatewayHost, gatewayPort, keyPath)
 		},
 	}
 
@@ -116,112 +103,24 @@ Example:
 	return cmd
 }
 
-// ensureSSHKey checks for an existing SSH key and generates one if missing.
-func ensureSSHKey() (string, error) {
-	homeDir, err := os.UserHomeDir()
-	if err != nil {
-		return "", err
-	}
-
-	candidates := []string{
-		filepath.Join(homeDir, ".ssh", "id_ed25519"),
-		filepath.Join(homeDir, ".ssh", "id_rsa"),
-	}
-	for _, p := range candidates {
-		if _, err := os.Stat(p); err == nil {
-			return p, nil
-		}
-	}
-
-	// Generate ed25519 key.
-	keyPath := candidates[0]
-	sshDir := filepath.Dir(keyPath)
-	if err := os.MkdirAll(sshDir, 0700); err != nil {
-		return "", fmt.Errorf("failed to create %s: %w", sshDir, err)
-	}
-
-	cmd := exec.Command("ssh-keygen", "-t", "ed25519", "-f", keyPath, "-N", "", "-q")
-	cmd.Stdin = os.Stdin
-	cmd.Stdout = os.Stderr
-	cmd.Stderr = os.Stderr
-	if err := cmd.Run(); err != nil {
-		return "", fmt.Errorf("ssh-keygen failed: %w", err)
-	}
-
-	return keyPath, nil
-}
-
-func sshConfigFilePath() (string, error) {
-	homeDir, err := os.UserHomeDir()
-	if err != nil {
-		return "", err
-	}
-	return filepath.Join(homeDir, ".ssh", "config"), nil
-}
-
-// buildSSHConfigEntry creates the SSH config block for a lakebox.
-// The lakebox ID is used as both the Host alias and the SSH User.
-func buildSSHConfigEntry(lakeboxID, host, port, keyPath string) string {
-	return fmt.Sprintf(`Host %s
-    HostName %s
-    Port %s
-    User %s
-    IdentityFile %s
-    IdentitiesOnly yes
-    PreferredAuthentications publickey
-    PasswordAuthentication no
-    KbdInteractiveAuthentication no
-    StrictHostKeyChecking no
-    UserKnownHostsFile /dev/null
-    LogLevel INFO
-`, lakeboxID, host, port, lakeboxID, keyPath)
-}
-
-// writeSSHConfigEntry idempotently writes a single lakebox entry to ~/.ssh/config.
-// Replaces any existing lakebox block in-place.
-func writeSSHConfigEntry(configPath, lakeboxID, entry string) error {
-	sshDir := filepath.Dir(configPath)
-	if err := os.MkdirAll(sshDir, 0700); err != nil {
-		return err
-	}
-
-	existing, err := os.ReadFile(configPath)
-	if err != nil && !os.IsNotExist(err) {
-		return err
-	}
-
-	wrappedEntry := fmt.Sprintf("%s\n%s%s\n", sshConfigMarkerStart, entry, sshConfigMarkerEnd)
-	content := string(existing)
-
-	// Remove existing lakebox block if present.
-	startIdx := strings.Index(content, sshConfigMarkerStart)
-	if startIdx >= 0 {
-		endIdx := strings.Index(content[startIdx:], sshConfigMarkerEnd)
-		if endIdx >= 0 {
-			endIdx += startIdx + len(sshConfigMarkerEnd)
-			if endIdx < len(content) && content[endIdx] == '\n' {
-				endIdx++
-			}
-			content = content[:startIdx] + content[endIdx:]
-		}
-	}
-
-	if !strings.HasSuffix(content, "\n") && len(content) > 0 {
-		content += "\n"
-	}
-	content += wrappedEntry
-
-	return os.WriteFile(configPath, []byte(content), 0600)
-}
-
-// execSSH execs into ssh using the lakebox ID as the Host alias.
-func execSSH(lakeboxID string) error {
+// execSSHDirect execs into ssh with all options passed as args (no ~/.ssh/config needed).
+func execSSHDirect(lakeboxID, host, port, keyPath string) error {
 	sshPath, err := exec.LookPath("ssh")
 	if err != nil {
 		return fmt.Errorf("ssh not found in PATH: %w", err)
 	}
 
-	args := []string{"ssh", lakeboxID}
+	args := []string{
+		"ssh",
+		"-i", keyPath,
+		"-p", port,
+		"-o", "IdentitiesOnly=yes",
+		"-o", "PreferredAuthentications=publickey",
+		"-o", "StrictHostKeyChecking=no",
+		"-o", "UserKnownHostsFile=/dev/null",
+		"-o", "LogLevel=ERROR",
+		fmt.Sprintf("%s@%s", lakeboxID, host),
+	}
 
 	if runtime.GOOS == "windows" {
 		cmd := exec.Command(sshPath, args[1:]...)

From 4b4186113ebf7cc8790ec9a0766ba2102d9f48cb Mon Sep 17 00:00:00 2001
From: shuochen0311 <shuo.chen@databricks.com>
Date: Tue, 14 Apr 2026 03:22:20 +0000
Subject: [PATCH 04/36] Auto-register SSH key after auth login, fix login hook
 matching

- Hook into auth login PostRun to auto-generate ~/.ssh/lakebox_rsa and
  register it after OAuth completes
- Fix hook: match on sub.Name() not sub.Use (Use includes args)
- Export EnsureAndReadKey and RegisterKey for use by auth hook
- Update help text

Co-authored-by: Isaac
---
 cmd/cmd.go              | 52 ++++++++++++++++++++++++++++++++++++++---
 cmd/lakebox/lakebox.go  |  3 +--
 cmd/lakebox/register.go | 23 ++++++++++++++++++
 3 files changed, 73 insertions(+), 5 deletions(-)

diff --git a/cmd/cmd.go b/cmd/cmd.go
index c120f25aa7..ddbb70f451 100644
--- a/cmd/cmd.go
+++ b/cmd/cmd.go
@@ -2,10 +2,12 @@ package cmd
 
 import (
 	"context"
+	"fmt"
 
 	"github.com/databricks/cli/cmd/auth"
 	"github.com/databricks/cli/cmd/lakebox"
 	"github.com/databricks/cli/cmd/root"
+	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/spf13/cobra"
 )
 
@@ -20,8 +22,7 @@ microVM isolation. Each lakebox is a personal sandbox with pre-installed
 tooling (Python, Node.js, Rust, Databricks CLI) and persistent storage.
 
 Getting started:
-  lakebox auth login --host https://...   # authenticate to Databricks
-  lakebox register                        # generate SSH key and register
+  lakebox auth login --host https://...   # authenticate to Databricks workspace and lakebox service
   lakebox ssh                             # SSH to your default lakebox
 
 Common workflows:
@@ -37,7 +38,52 @@ The CLI manages your ~/.ssh/config so you can also connect directly:
 `
 	cli.CompletionOptions.DisableDefaultCmd = true
 
-	cli.AddCommand(auth.New())
+	authCmd := auth.New()
+	// Hook into 'auth login' to auto-register SSH key after OAuth completes.
+	for _, sub := range authCmd.Commands() {
+		if sub.Name() == "login" {
+			origRunE := sub.RunE
+			sub.RunE = func(cmd *cobra.Command, args []string) error {
+				// Run the original auth login.
+				if err := origRunE(cmd, args); err != nil {
+					return err
+				}
+
+				// Auto-register: generate lakebox SSH key and register it.
+				fmt.Fprintln(cmd.ErrOrStderr(), "")
+				fmt.Fprintln(cmd.ErrOrStderr(), "Setting up SSH access...")
+
+				keyPath, pubKey, err := lakebox.EnsureAndReadKey()
+				if err != nil {
+					fmt.Fprintf(cmd.ErrOrStderr(),
+						"SSH key setup failed: %v\n"+
+							"You can set it up later with: lakebox register\n", err)
+					return nil
+				}
+				fmt.Fprintf(cmd.ErrOrStderr(), "Using SSH key: %s\n", keyPath)
+
+				if err := root.MustWorkspaceClient(cmd, args); err != nil {
+					fmt.Fprintf(cmd.ErrOrStderr(),
+						"Could not initialize workspace client for key registration.\n"+
+							"Run 'lakebox register' to complete setup.\n")
+					return nil
+				}
+
+				w := cmdctx.WorkspaceClient(cmd.Context())
+				if err := lakebox.RegisterKey(cmd.Context(), w, pubKey); err != nil {
+					fmt.Fprintf(cmd.ErrOrStderr(),
+						"Key registration failed: %v\n"+
+							"Run 'lakebox register' to retry.\n", err)
+					return nil
+				}
+
+				fmt.Fprintln(cmd.ErrOrStderr(), "SSH key registered. You're ready to use 'lakebox ssh'.")
+				return nil
+			}
+			break
+		}
+	}
+	cli.AddCommand(authCmd)
 
 	// Register lakebox subcommands directly at root level.
 	for _, sub := range lakebox.New().Commands() {
diff --git a/cmd/lakebox/lakebox.go b/cmd/lakebox/lakebox.go
index 127b5d93bf..4afa321241 100644
--- a/cmd/lakebox/lakebox.go
+++ b/cmd/lakebox/lakebox.go
@@ -15,8 +15,7 @@ microVM isolation. Each lakebox is a personal sandbox with pre-installed
 tooling (Python, Node.js, Rust, Databricks CLI) and persistent storage.
 
 Getting started:
-  lakebox auth login --host https://...   # authenticate to Databricks
-  lakebox register                        # generate SSH key and register
+  lakebox auth login --host https://...   # authenticate to Databricks workspace and lakebox service
   lakebox ssh                             # SSH to your default lakebox
 
 Common workflows:
diff --git a/cmd/lakebox/register.go b/cmd/lakebox/register.go
index 7286a14bf5..a1da60422b 100644
--- a/cmd/lakebox/register.go
+++ b/cmd/lakebox/register.go
@@ -1,6 +1,7 @@
 package lakebox
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"os/exec"
@@ -8,6 +9,7 @@ import (
 
 	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdctx"
+	"github.com/databricks/databricks-sdk-go"
 	"github.com/spf13/cobra"
 )
 
@@ -108,3 +110,24 @@ func ensureLakeboxKey() (string, bool, error) {
 
 	return keyPath, true, nil
 }
+
+// EnsureAndReadKey generates the lakebox SSH key if needed and returns
+// (keyPath, publicKeyContent, error). Exported for use by the auth login hook.
+func EnsureAndReadKey() (string, string, error) {
+	keyPath, _, err := ensureLakeboxKey()
+	if err != nil {
+		return "", "", err
+	}
+	pubKeyData, err := os.ReadFile(keyPath + ".pub")
+	if err != nil {
+		return "", "", fmt.Errorf("failed to read public key %s.pub: %w", keyPath, err)
+	}
+	return keyPath, string(pubKeyData), nil
+}
+
+// RegisterKey registers a public key with the lakebox API. Exported for use
+// by the auth login hook.
+func RegisterKey(ctx context.Context, w *databricks.WorkspaceClient, pubKey string) error {
+	api := newLakeboxAPI(w)
+	return api.registerKey(ctx, pubKey)
+}

From df599e9273beebe13ec81c7d5341b4b77001e2a8 Mon Sep 17 00:00:00 2001
From: shuochen0311 <shuo.chen@databricks.com>
Date: Tue, 14 Apr 2026 21:42:07 +0000
Subject: [PATCH 05/36] Support passthrough args and remote commands in lakebox
 ssh

Everything after -- is passed directly to the ssh process, enabling:
  lakebox ssh -- echo hello          # run command and return
  lakebox ssh <id> -- cat /etc/os-release
  lakebox ssh -- -L 8080:localhost:8080  # port forwarding

Co-authored-by: Isaac
---
 cmd/lakebox/ssh.go | 90 +++++++++++++++++++++++++---------------------
 1 file changed, 50 insertions(+), 40 deletions(-)

diff --git a/cmd/lakebox/ssh.go b/cmd/lakebox/ssh.go
index 8868f38e81..7559893bfb 100644
--- a/cmd/lakebox/ssh.go
+++ b/cmd/lakebox/ssh.go
@@ -21,24 +21,24 @@ func newSSHCommand() *cobra.Command {
 	var gatewayPort string
 
 	cmd := &cobra.Command{
-		Use:   "ssh [lakebox-id]",
+		Use:   "ssh [lakebox-id] [-- <ssh-args-or-command>...]",
 		Short: "SSH into a Lakebox environment",
 		Long: `SSH into a Lakebox environment.
 
-This command:
-1. Authenticates to the Databricks workspace
-2. Ensures you have a local SSH key (~/.ssh/id_ed25519)
-3. Creates a lakebox if one doesn't exist (installs your public key)
-4. Updates ~/.ssh/config with a Host entry for the lakebox
-5. Connects via SSH using the lakebox ID as the SSH username
-
-Without arguments, creates a new lakebox. With a lakebox ID argument,
-connects to the specified lakebox.
-
-Example:
-  databricks lakebox ssh                    # create and connect to a new lakebox
-  databricks lakebox ssh happy-panda-1234   # connect to existing lakebox`,
-		Args: cobra.MaximumNArgs(1),
+Connect to your default or a named lakebox via SSH. Extra arguments
+after -- are passed directly to the ssh process. This lets you run
+remote commands, set up port forwarding, or pass any other ssh flags.
+
+Examples:
+  lakebox ssh                                  # interactive shell on default lakebox
+  lakebox ssh happy-panda-1234                 # interactive shell on specific lakebox
+  lakebox ssh -- ls -la /home                  # run command on default lakebox
+  lakebox ssh happy-panda-1234 -- cat /etc/os-release  # run command on specific lakebox
+  lakebox ssh -- -L 8080:localhost:8080        # port forwarding on default lakebox`,
+		// Disable flag parsing after -- so extra args are passed through.
+		DisableFlagParsing: false,
+		// Accept any number of args: [lakebox-id] [-- extra...]
+		Args: cobra.ArbitraryArgs,
 		PreRunE: func(cmd *cobra.Command, args []string) error {
 			return root.MustWorkspaceClient(cmd, args)
 		},
@@ -61,39 +61,47 @@ Example:
 			}
 			fmt.Fprintf(cmd.ErrOrStderr(), "Using SSH key: %s\n", keyPath)
 
-			// Determine lakebox ID:
-			// 1. Explicit arg → use it
-			// 2. Local default exists → use it
-			// 3. Neither → create a new one and set as default
+			// Parse args: first arg (if not starting with -) is lakebox ID,
+			// everything else is passed through to ssh.
 			var lakeboxID string
-			if len(args) > 0 {
+			var extraArgs []string
+
+			if len(args) > 0 && args[0] != "--" && args[0][0] != '-' {
 				lakeboxID = args[0]
-			} else if def := getDefault(profile); def != "" {
-				lakeboxID = def
-				fmt.Fprintf(cmd.ErrOrStderr(), "Using default lakebox: %s\n", lakeboxID)
+				extraArgs = args[1:]
 			} else {
-				api := newLakeboxAPI(w)
-				pubKeyData, err := os.ReadFile(keyPath + ".pub")
-				if err != nil {
-					return fmt.Errorf("failed to read public key %s.pub: %w", keyPath, err)
-				}
-
-				fmt.Fprintf(cmd.ErrOrStderr(), "Creating lakebox...\n")
-				result, err := api.create(ctx, string(pubKeyData))
-				if err != nil {
-					return fmt.Errorf("failed to create lakebox: %w", err)
-				}
-				lakeboxID = result.LakeboxID
-				fmt.Fprintf(cmd.ErrOrStderr(), "Lakebox %s created (status: %s)\n", lakeboxID, result.Status)
+				extraArgs = args
+			}
 
-				if err := setDefault(profile, lakeboxID); err != nil {
-					fmt.Fprintf(cmd.ErrOrStderr(), "Warning: failed to save default: %v\n", err)
+			// Determine lakebox ID if not explicit.
+			if lakeboxID == "" {
+				if def := getDefault(profile); def != "" {
+					lakeboxID = def
+					fmt.Fprintf(cmd.ErrOrStderr(), "Using default lakebox: %s\n", lakeboxID)
+				} else {
+					api := newLakeboxAPI(w)
+					pubKeyData, err := os.ReadFile(keyPath + ".pub")
+					if err != nil {
+						return fmt.Errorf("failed to read public key %s.pub: %w", keyPath, err)
+					}
+
+					fmt.Fprintf(cmd.ErrOrStderr(), "Creating lakebox...\n")
+					result, err := api.create(ctx, string(pubKeyData))
+					if err != nil {
+						return fmt.Errorf("failed to create lakebox: %w", err)
+					}
+					lakeboxID = result.LakeboxID
+					fmt.Fprintf(cmd.ErrOrStderr(), "Lakebox %s created (status: %s)\n", lakeboxID, result.Status)
+
+					if err := setDefault(profile, lakeboxID); err != nil {
+						fmt.Fprintf(cmd.ErrOrStderr(), "Warning: failed to save default: %v\n", err)
+					}
 				}
 			}
 
 			fmt.Fprintf(cmd.ErrOrStderr(), "Connecting to %s@%s:%s...\n",
 				lakeboxID, gatewayHost, gatewayPort)
-			return execSSHDirect(lakeboxID, gatewayHost, gatewayPort, keyPath)
+			return execSSHDirect(lakeboxID, gatewayHost, gatewayPort, keyPath, extraArgs)
 		},
 	}
 
@@ -104,7 +112,8 @@ Example:
 }
 
 // execSSHDirect execs into ssh with all options passed as args (no ~/.ssh/config needed).
-func execSSHDirect(lakeboxID, host, port, keyPath string) error {
+// Extra args are appended after the destination (for remote commands or ssh flags).
+func execSSHDirect(lakeboxID, host, port, keyPath string, extraArgs []string) error {
 	sshPath, err := exec.LookPath("ssh")
 	if err != nil {
 		return fmt.Errorf("ssh not found in PATH: %w", err)
@@ -121,6 +130,7 @@ func execSSHDirect(lakeboxID, host, port, keyPath string) error {
 		"-o", "LogLevel=ERROR",
 		fmt.Sprintf("%s@%s", lakeboxID, host),
 	}
+	args = append(args, extraArgs...)
 
 	if runtime.GOOS == "windows" {
 		cmd := exec.Command(sshPath, args[1:]...)

From cd2579760e15edbc4af014da6bc9963a4669110e Mon Sep 17 00:00:00 2001
From: Stas Kelvich <stas.kelvich@databricks.com>
Date: Tue, 14 Apr 2026 15:30:12 -0700
Subject: [PATCH 06/36] Fix workspace client init after login, persist last
 profile

After 'lakebox auth login --host <url>', the post-login hook now
constructs the workspace client directly from the --host/--profile
flags instead of using MustWorkspaceClient (which started with an
empty config and fell back to the DEFAULT profile).

All lakebox commands now use a mustWorkspaceClient wrapper that reads
the last-login profile from ~/.databricks/lakebox.json, so 'lakebox ssh'
uses the correct profile without requiring --profile on every invocation.

Also adds install.sh and upload.sh scripts.
---
 cmd/cmd.go              | 30 +++++++++++++---
 cmd/lakebox/create.go   |  3 +-
 cmd/lakebox/default.go  |  3 +-
 cmd/lakebox/delete.go   |  3 +-
 cmd/lakebox/lakebox.go  | 15 +++++++-
 cmd/lakebox/list.go     |  3 +-
 cmd/lakebox/register.go |  3 +-
 cmd/lakebox/ssh.go      |  5 +--
 cmd/lakebox/state.go    | 21 +++++++++++
 cmd/lakebox/status.go   |  3 +-
 install.sh              | 80 +++++++++++++++++++++++++++++++++++++++++
 upload.sh               | 13 +++++++
 12 files changed, 160 insertions(+), 22 deletions(-)
 create mode 100755 install.sh
 create mode 100755 upload.sh

diff --git a/cmd/cmd.go b/cmd/cmd.go
index ddbb70f451..8a70375514 100644
--- a/cmd/cmd.go
+++ b/cmd/cmd.go
@@ -3,11 +3,12 @@ package cmd
 import (
 	"context"
 	"fmt"
+	"strings"
 
 	"github.com/databricks/cli/cmd/auth"
 	"github.com/databricks/cli/cmd/lakebox"
 	"github.com/databricks/cli/cmd/root"
-	"github.com/databricks/cli/libs/cmdctx"
+	"github.com/databricks/databricks-sdk-go"
 	"github.com/spf13/cobra"
 )
 
@@ -62,14 +63,33 @@ The CLI manages your ~/.ssh/config so you can also connect directly:
 				}
 				fmt.Fprintf(cmd.ErrOrStderr(), "Using SSH key: %s\n", keyPath)
 
-				if err := root.MustWorkspaceClient(cmd, args); err != nil {
+				host := cmd.Flag("host").Value.String()
+				if host == "" && len(args) > 0 {
+					host = args[0]
+				}
+				profile := cmd.Flag("profile").Value.String()
+				if profile == "" && host != "" {
+					// Derive profile name the same way auth login does.
+					h := strings.TrimPrefix(host, "https://")
+					h = strings.TrimPrefix(h, "http://")
+					profile = strings.SplitN(h, ".", 2)[0]
+				}
+				if profile != "" {
+					if err := lakebox.SetLastProfile(profile); err != nil {
+						fmt.Fprintf(cmd.ErrOrStderr(), "Warning: failed to save last profile: %v\n", err)
+					}
+				}
+				w, err := databricks.NewWorkspaceClient(&databricks.Config{
+					Host:    host,
+					Profile: profile,
+				})
+				if err != nil {
 					fmt.Fprintf(cmd.ErrOrStderr(),
-						"Could not initialize workspace client for key registration.\n"+
-							"Run 'lakebox register' to complete setup.\n")
+						"Could not initialize workspace client for key registration: %v\n"+
+							"Run 'lakebox register' to complete setup.\n", err)
 					return nil
 				}
 
-				w := cmdctx.WorkspaceClient(cmd.Context())
 				if err := lakebox.RegisterKey(cmd.Context(), w, pubKey); err != nil {
 					fmt.Fprintf(cmd.ErrOrStderr(),
 						"Key registration failed: %v\n"+
diff --git a/cmd/lakebox/create.go b/cmd/lakebox/create.go
index 872776cc8d..db1a22ebb7 100644
--- a/cmd/lakebox/create.go
+++ b/cmd/lakebox/create.go
@@ -4,7 +4,6 @@ import (
 	"fmt"
 	"os"
 
-	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/spf13/cobra"
 )
@@ -26,7 +25,7 @@ authorized_keys so you can SSH directly. Otherwise the gateway key is used.
 Example:
   databricks lakebox create
   databricks lakebox create --public-key-file ~/.ssh/id_ed25519.pub`,
-		PreRunE: root.MustWorkspaceClient,
+		PreRunE: mustWorkspaceClient,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := cmd.Context()
 			w := cmdctx.WorkspaceClient(ctx)
diff --git a/cmd/lakebox/default.go b/cmd/lakebox/default.go
index 9d5a366c9c..b632c5984a 100644
--- a/cmd/lakebox/default.go
+++ b/cmd/lakebox/default.go
@@ -3,7 +3,6 @@ package lakebox
 import (
 	"fmt"
 
-	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/spf13/cobra"
 )
@@ -19,7 +18,7 @@ The default is stored locally in ~/.databricks/lakebox.json per profile.
 Example:
   databricks lakebox set-default happy-panda-1234`,
 		Args:    cobra.ExactArgs(1),
-		PreRunE: root.MustWorkspaceClient,
+		PreRunE: mustWorkspaceClient,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			w := cmdctx.WorkspaceClient(cmd.Context())
 			profile := w.Config.Profile
diff --git a/cmd/lakebox/delete.go b/cmd/lakebox/delete.go
index a814083ed3..9c8ce93963 100644
--- a/cmd/lakebox/delete.go
+++ b/cmd/lakebox/delete.go
@@ -3,7 +3,6 @@ package lakebox
 import (
 	"fmt"
 
-	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/spf13/cobra"
 )
@@ -20,7 +19,7 @@ creator (same auth token) can delete a lakebox.
 Example:
   databricks lakebox delete happy-panda-1234`,
 		Args:    cobra.ExactArgs(1),
-		PreRunE: root.MustWorkspaceClient,
+		PreRunE: mustWorkspaceClient,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := cmd.Context()
 			w := cmdctx.WorkspaceClient(ctx)
diff --git a/cmd/lakebox/lakebox.go b/cmd/lakebox/lakebox.go
index 4afa321241..6a968df87a 100644
--- a/cmd/lakebox/lakebox.go
+++ b/cmd/lakebox/lakebox.go
@@ -1,6 +1,7 @@
 package lakebox
 
 import (
+	"github.com/databricks/cli/cmd/root"
 	"github.com/spf13/cobra"
 )
 
@@ -32,12 +33,24 @@ The CLI manages your ~/.ssh/config so you can also connect directly:
 	}
 
 	cmd.AddCommand(newRegisterCommand())
+	cmd.AddCommand(newSetDefaultCommand())
 	cmd.AddCommand(newSSHCommand())
 	cmd.AddCommand(newListCommand())
 	cmd.AddCommand(newCreateCommand())
 	cmd.AddCommand(newDeleteCommand())
 	cmd.AddCommand(newStatusCommand())
-	cmd.AddCommand(newSetDefaultCommand())
 
 	return cmd
 }
+
+// mustWorkspaceClient applies the saved last-login profile when the user
+// hasn't explicitly set --profile, then delegates to root.MustWorkspaceClient.
+func mustWorkspaceClient(cmd *cobra.Command, args []string) error {
+	profileFlag := cmd.Flag("profile")
+	if profileFlag != nil && !profileFlag.Changed {
+		if last := GetLastProfile(); last != "" {
+			_ = profileFlag.Value.Set(last)
+		}
+	}
+	return root.MustWorkspaceClient(cmd, args)
+}
diff --git a/cmd/lakebox/list.go b/cmd/lakebox/list.go
index 90139d6be8..3222d1c10c 100644
--- a/cmd/lakebox/list.go
+++ b/cmd/lakebox/list.go
@@ -4,7 +4,6 @@ import (
 	"encoding/json"
 	"fmt"
 
-	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/spf13/cobra"
 )
@@ -23,7 +22,7 @@ current status and ID.
 Example:
   databricks lakebox list
   databricks lakebox list --json`,
-		PreRunE: root.MustWorkspaceClient,
+		PreRunE: mustWorkspaceClient,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := cmd.Context()
 			w := cmdctx.WorkspaceClient(ctx)
diff --git a/cmd/lakebox/register.go b/cmd/lakebox/register.go
index a1da60422b..27d6cc59a1 100644
--- a/cmd/lakebox/register.go
+++ b/cmd/lakebox/register.go
@@ -7,7 +7,6 @@ import (
 	"os/exec"
 	"path/filepath"
 
-	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/databricks/databricks-sdk-go"
 	"github.com/spf13/cobra"
@@ -30,7 +29,7 @@ Run this once per machine.
 
 Example:
   lakebox register`,
-		PreRunE: root.MustWorkspaceClient,
+		PreRunE: mustWorkspaceClient,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := cmd.Context()
 			w := cmdctx.WorkspaceClient(ctx)
diff --git a/cmd/lakebox/ssh.go b/cmd/lakebox/ssh.go
index 8868f38e81..86098baf5a 100644
--- a/cmd/lakebox/ssh.go
+++ b/cmd/lakebox/ssh.go
@@ -6,7 +6,6 @@ import (
 	"os/exec"
 	"runtime"
 
-	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/spf13/cobra"
 )
@@ -39,9 +38,7 @@ Example:
   databricks lakebox ssh                    # create and connect to a new lakebox
   databricks lakebox ssh happy-panda-1234   # connect to existing lakebox`,
 		Args: cobra.MaximumNArgs(1),
-		PreRunE: func(cmd *cobra.Command, args []string) error {
-			return root.MustWorkspaceClient(cmd, args)
-		},
+		PreRunE: mustWorkspaceClient,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := cmd.Context()
 			w := cmdctx.WorkspaceClient(ctx)
diff --git a/cmd/lakebox/state.go b/cmd/lakebox/state.go
index c0c8ad2d84..b84b5b16e1 100644
--- a/cmd/lakebox/state.go
+++ b/cmd/lakebox/state.go
@@ -12,6 +12,8 @@ import (
 type stateFile struct {
 	// Profile name → default lakebox ID.
 	Defaults map[string]string `json:"defaults"`
+	// Last profile used with 'lakebox auth login'.
+	LastProfile string `json:"last_profile,omitempty"`
 }
 
 func stateFilePath() (string, error) {
@@ -80,6 +82,25 @@ func setDefault(profile, lakeboxID string) error {
 	return saveState(state)
 }
 
+// GetLastProfile returns the profile saved by the most recent 'lakebox auth login'.
+func GetLastProfile() string {
+	state, err := loadState()
+	if err != nil {
+		return ""
+	}
+	return state.LastProfile
+}
+
+// SetLastProfile persists the profile used during 'lakebox auth login'.
+func SetLastProfile(profile string) error {
+	state, err := loadState()
+	if err != nil {
+		return err
+	}
+	state.LastProfile = profile
+	return saveState(state)
+}
+
 func clearDefault(profile string) error {
 	state, err := loadState()
 	if err != nil {
diff --git a/cmd/lakebox/status.go b/cmd/lakebox/status.go
index 4bb130496d..eaeeb8d7cc 100644
--- a/cmd/lakebox/status.go
+++ b/cmd/lakebox/status.go
@@ -4,7 +4,6 @@ import (
 	"encoding/json"
 	"fmt"
 
-	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/spf13/cobra"
 )
@@ -21,7 +20,7 @@ Example:
   databricks lakebox status happy-panda-1234
   databricks lakebox status happy-panda-1234 --json`,
 		Args:    cobra.ExactArgs(1),
-		PreRunE: root.MustWorkspaceClient,
+		PreRunE: mustWorkspaceClient,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := cmd.Context()
 			w := cmdctx.WorkspaceClient(ctx)
diff --git a/install.sh b/install.sh
new file mode 100755
index 0000000000..acdf259b4c
--- /dev/null
+++ b/install.sh
@@ -0,0 +1,80 @@
+#!/bin/sh
+# Lakebox CLI installer — . <(curl -s devbox.dbrx.dev)
+
+_lakebox_install() {
+    INSTALL_DIR="$HOME/.lakebox/bin"
+    REMOTE_NAME="databricks"
+    LOCAL_NAME="lakebox"
+    BASE_URL="https://devbox.dbrx.dev"
+
+    case "$(uname -s)" in
+        Linux*)  OS="linux" ;;
+        Darwin*) OS="darwin" ;;
+        *)  printf "error: unsupported OS: %s\n" "$(uname -s)" >&2; return 1 ;;
+    esac
+
+    case "$(uname -m)" in
+        x86_64|amd64)  ARCH="amd64" ;;
+        aarch64|arm64) ARCH="arm64" ;;
+        *)  printf "error: unsupported arch: %s\n" "$(uname -m)" >&2; return 1 ;;
+    esac
+
+    url="${BASE_URL}/${REMOTE_NAME}-${OS}-${ARCH}"
+
+    printf "📦 Installing Lakebox CLI (%s/%s)...\n" "$OS" "$ARCH"
+
+    mkdir -p "$INSTALL_DIR" || { printf "error: could not create %s\n" "$INSTALL_DIR" >&2; return 1; }
+
+    if command -v curl >/dev/null 2>&1; then
+        curl -fSL --progress-bar "$url" -o "$INSTALL_DIR/$LOCAL_NAME" || { printf "error: download failed\n" >&2; return 1; }
+    elif command -v wget >/dev/null 2>&1; then
+        wget -q --show-progress "$url" -O "$INSTALL_DIR/$LOCAL_NAME" || { printf "error: download failed\n" >&2; return 1; }
+    else
+        printf "error: curl or wget is required\n" >&2; return 1
+    fi
+
+    chmod +x "$INSTALL_DIR/$LOCAL_NAME"
+
+    PATH_LINE="export PATH=\"\$HOME/.lakebox/bin:\$PATH\""
+    case ":$PATH:" in
+        *":$INSTALL_DIR:"*) ;;
+        *)
+            added=0
+            for rc in "$HOME/.zshrc" "$HOME/.bashrc"; do
+                [ -f "$rc" ] || continue
+                if ! grep -qF '.lakebox/bin' "$rc" 2>/dev/null; then
+                    printf '\n# Lakebox CLI\n%s\n' "$PATH_LINE" >> "$rc"
+                    printf "📝 Updated %s\n" "$rc"
+                    added=1
+                fi
+            done
+            if [ "$added" = 0 ]; then
+                if [ "$OS" = "darwin" ]; then
+                    rc="$HOME/.zshrc"
+                else
+                    rc="$HOME/.bashrc"
+                fi
+                printf '\n# Lakebox CLI\n%s\n' "$PATH_LINE" >> "$rc"
+                printf "📝 Updated %s\n" "$rc"
+            fi
+            export PATH="$INSTALL_DIR:$PATH"
+            ;;
+    esac
+
+    printf "\n✅ Lakebox CLI installed to %s\n" "$INSTALL_DIR/$LOCAL_NAME"
+
+    LAKEBOX_HOST="https://dbsql-dev-testing-default.dev.databricks.com"
+    LAKEBOX_PROFILE="dbsql-dev-testing-default"
+    if ! grep -qF "$LAKEBOX_PROFILE" "$HOME/.databrickscfg" 2>/dev/null; then
+        printf "\n🔑 Logging in...\n"
+        lakebox auth login --host "$LAKEBOX_HOST" --profile "$LAKEBOX_PROFILE"
+    fi
+
+    printf "\nCommon workflows:\n"
+    printf "  lakebox ssh                             # SSH to your default lakebox\n"
+    printf "  lakebox ssh my-project                  # SSH to a named lakebox\n"
+    printf "  lakebox list                            # list your lakeboxes\n"
+}
+
+_lakebox_install
+unset -f _lakebox_install
\ No newline at end of file
diff --git a/upload.sh b/upload.sh
new file mode 100755
index 0000000000..c55c0aa182
--- /dev/null
+++ b/upload.sh
@@ -0,0 +1,13 @@
+#!/bin/sh
+set -eu
+
+HOST="arca.ssh"
+FILES="install.sh databricks-darwin-amd64 databricks-darwin-arm64 databricks-linux-amd64 databricks-linux-arm64"
+
+for f in $FILES; do
+    printf "Uploading %s...\n" "$f"
+    scp "$f" "$HOST:~/"
+    ssh "$HOST" "~/unp-upload.sh ~/$f"
+done
+
+printf "\nDone.\n"

From c1168a414f429f47517557ab8b9ce3bc7084fe28 Mon Sep 17 00:00:00 2001
From: shuochen0311 <shuo.chen@databricks.com>
Date: Thu, 16 Apr 2026 06:48:54 +0000
Subject: [PATCH 07/36] Add consistent terminal UI: spinners, colors, aligned
 output

Single cyan accent color throughout. Bold for IDs, dim for metadata.
Braille spinner with elapsed time during async operations.

- create: animated spinner during provisioning
- list: aligned columns with colored status, cyan bold for running
- status: clean field layout
- delete: spinner during removal
- ssh: spinner during connection
- register: spinner during key registration
- Shared ui.go with all primitives

Co-authored-by: Isaac
---
 cmd/lakebox/create.go   |  21 +++---
 cmd/lakebox/delete.go   |  15 +++--
 cmd/lakebox/list.go     |  46 +++++++++++--
 cmd/lakebox/register.go |  11 +++-
 cmd/lakebox/ssh.go      |  17 ++---
 cmd/lakebox/status.go   |  13 ++--
 cmd/lakebox/ui.go       | 141 ++++++++++++++++++++++++++++++++++++++++
 7 files changed, 221 insertions(+), 43 deletions(-)
 create mode 100644 cmd/lakebox/ui.go

diff --git a/cmd/lakebox/create.go b/cmd/lakebox/create.go
index db1a22ebb7..c4ce3a439e 100644
--- a/cmd/lakebox/create.go
+++ b/cmd/lakebox/create.go
@@ -19,17 +19,14 @@ func newCreateCommand() *cobra.Command {
 Creates a new personal development environment backed by a microVM.
 Blocks until the lakebox is running and prints the lakebox ID.
 
-If --public-key-file is provided, the key is installed in the lakebox's
-authorized_keys so you can SSH directly. Otherwise the gateway key is used.
-
 Example:
-  databricks lakebox create
-  databricks lakebox create --public-key-file ~/.ssh/id_ed25519.pub`,
+  lakebox create`,
 		PreRunE: mustWorkspaceClient,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := cmd.Context()
 			w := cmdctx.WorkspaceClient(ctx)
 			api := newLakeboxAPI(w)
+			stderr := cmd.ErrOrStderr()
 
 			var publicKey string
 			if publicKeyFile != "" {
@@ -40,37 +37,37 @@ Example:
 				publicKey = string(data)
 			}
 
-			fmt.Fprintf(cmd.ErrOrStderr(), "Creating lakebox...\n")
+			s := spin(stderr, "Provisioning your lakebox…")
 
 			result, err := api.create(ctx, publicKey)
 			if err != nil {
+				s.fail("Failed to create lakebox")
 				return fmt.Errorf("failed to create lakebox: %w", err)
 			}
 
+			s.ok(fmt.Sprintf("Lakebox %s is %s", bold(result.LakeboxID), status(result.Status)))
+
 			profile := w.Config.Profile
 			if profile == "" {
 				profile = w.Config.Host
 			}
 
-			// Set as default if no default exists, or the current default
-			// has been deleted (no longer in the list).
 			currentDefault := getDefault(profile)
 			shouldSetDefault := currentDefault == ""
 			if !shouldSetDefault && currentDefault != "" {
-				// Check if the current default still exists.
 				if _, err := api.get(ctx, currentDefault); err != nil {
 					shouldSetDefault = true
 				}
 			}
 			if shouldSetDefault {
 				if err := setDefault(profile, result.LakeboxID); err != nil {
-					fmt.Fprintf(cmd.ErrOrStderr(), "Warning: failed to save default: %v\n", err)
+					warn(stderr, fmt.Sprintf("Could not save default: %v", err))
 				} else {
-					fmt.Fprintf(cmd.ErrOrStderr(), "Set as default lakebox.\n")
+					field(stderr, "default", result.LakeboxID)
 				}
 			}
 
-			fmt.Fprintf(cmd.ErrOrStderr(), "Lakebox created (status: %s)\n", result.Status)
+			blank(stderr)
 			fmt.Fprintln(cmd.OutOrStdout(), result.LakeboxID)
 			return nil
 		},
diff --git a/cmd/lakebox/delete.go b/cmd/lakebox/delete.go
index 9c8ce93963..ba56e2a508 100644
--- a/cmd/lakebox/delete.go
+++ b/cmd/lakebox/delete.go
@@ -13,35 +13,36 @@ func newDeleteCommand() *cobra.Command {
 		Short: "Delete a Lakebox environment",
 		Long: `Delete a Lakebox environment.
 
-Permanently terminates and removes the specified lakebox. Only the
-creator (same auth token) can delete a lakebox.
+Permanently terminates and removes the specified lakebox.
 
 Example:
-  databricks lakebox delete happy-panda-1234`,
+  lakebox delete happy-panda-1234`,
 		Args:    cobra.ExactArgs(1),
 		PreRunE: mustWorkspaceClient,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := cmd.Context()
 			w := cmdctx.WorkspaceClient(ctx)
 			api := newLakeboxAPI(w)
+			stderr := cmd.ErrOrStderr()
 
 			lakeboxID := args[0]
+			s := spin(stderr, fmt.Sprintf("Removing %s…", lakeboxID))
 
 			if err := api.delete(ctx, lakeboxID); err != nil {
+				s.fail(fmt.Sprintf("Failed to delete %s", lakeboxID))
 				return fmt.Errorf("failed to delete lakebox %s: %w", lakeboxID, err)
 			}
 
-			// Clear default if we just deleted it.
 			profile := w.Config.Profile
 			if profile == "" {
 				profile = w.Config.Host
 			}
 			if getDefault(profile) == lakeboxID {
 				_ = clearDefault(profile)
-				fmt.Fprintf(cmd.ErrOrStderr(), "Cleared default lakebox.\n")
+				s.ok(fmt.Sprintf("Removed %s %s", bold(lakeboxID), dim("(default cleared)")))
+			} else {
+				s.ok(fmt.Sprintf("Removed %s", bold(lakeboxID)))
 			}
-
-			fmt.Fprintf(cmd.ErrOrStderr(), "Deleted lakebox %s\n", lakeboxID)
 			return nil
 		},
 	}
diff --git a/cmd/lakebox/list.go b/cmd/lakebox/list.go
index 3222d1c10c..2ed3149658 100644
--- a/cmd/lakebox/list.go
+++ b/cmd/lakebox/list.go
@@ -3,6 +3,7 @@ package lakebox
 import (
 	"encoding/json"
 	"fmt"
+	"strings"
 
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/spf13/cobra"
@@ -20,8 +21,8 @@ Shows all lakeboxes associated with your account, including their
 current status and ID.
 
 Example:
-  databricks lakebox list
-  databricks lakebox list --json`,
+  lakebox list
+  lakebox list --json`,
 		PreRunE: mustWorkspaceClient,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := cmd.Context()
@@ -40,7 +41,7 @@ Example:
 			}
 
 			if len(entries) == 0 {
-				fmt.Fprintln(cmd.ErrOrStderr(), "No lakeboxes found.")
+				fmt.Fprintf(cmd.ErrOrStderr(), "  %sNo lakeboxes found.%s\n", dm, rs)
 				return nil
 			}
 
@@ -50,15 +51,48 @@ Example:
 			}
 			defaultID := getDefault(profile)
 
-			fmt.Fprintf(cmd.OutOrStdout(), "  %-30s %-12s %s\n", "ID", "STATUS", "DEFAULT")
+			out := cmd.OutOrStdout()
+
+			// Compute column width.
+			col := 10
+			for _, e := range entries {
+				if l := len(extractLakeboxID(e.Name)); l > col {
+					col = l
+				}
+			}
+			col += 2
+
+			blank(out)
+			fmt.Fprintf(out, "  %s%-*s  %-10s  %s%s\n", dm, col, "ID", "STATUS", "DEFAULT", rs)
+			fmt.Fprintf(out, "  %s%s%s\n", dm, strings.Repeat("─", col+22), rs)
+
 			for _, e := range entries {
 				id := extractLakeboxID(e.Name)
 				def := ""
 				if id == defaultID {
-					def = "*"
+					def = accent("*")
+				}
+				// Pad ID manually to avoid ANSI codes breaking alignment.
+				idPad := col - len(id)
+				if idPad < 0 {
+					idPad = 0
+				}
+				st := status(e.Status)
+				// Pad status to 10 visible chars.
+				stPad := 10 - len(e.Status)
+				if stPad < 0 {
+					stPad = 0
+				}
+				idStr := bold(id)
+				if strings.EqualFold(e.Status, "running") {
+					idStr = cyan + bo + id + rs
 				}
-				fmt.Fprintf(cmd.OutOrStdout(), "  %-30s %-12s %s\n", id, e.Status, def)
+				fmt.Fprintf(out, "  %s%s  %s%s  %s\n",
+					idStr, strings.Repeat(" ", idPad),
+					st, strings.Repeat(" ", stPad),
+					def)
 			}
+			blank(out)
 			return nil
 		},
 	}
diff --git a/cmd/lakebox/register.go b/cmd/lakebox/register.go
index 27d6cc59a1..f3550d8e5d 100644
--- a/cmd/lakebox/register.go
+++ b/cmd/lakebox/register.go
@@ -40,10 +40,11 @@ Example:
 				return fmt.Errorf("failed to ensure lakebox SSH key: %w", err)
 			}
 
+			stderr := cmd.ErrOrStderr()
 			if generated {
-				fmt.Fprintf(cmd.ErrOrStderr(), "Generated SSH key: %s\n", keyPath)
+				ok(stderr, fmt.Sprintf("Generated SSH key at %s", dim(keyPath)))
 			} else {
-				fmt.Fprintf(cmd.ErrOrStderr(), "Using existing SSH key: %s\n", keyPath)
+				field(stderr, "key", keyPath)
 			}
 
 			pubKeyData, err := os.ReadFile(keyPath + ".pub")
@@ -51,11 +52,15 @@ Example:
 				return fmt.Errorf("failed to read public key %s.pub: %w", keyPath, err)
 			}
 
+			s := spin(stderr, "Registering key…")
 			if err := api.registerKey(ctx, string(pubKeyData)); err != nil {
+				s.fail("Failed to register key")
 				return fmt.Errorf("failed to register key: %w", err)
 			}
+			s.ok("SSH key registered")
 
-			fmt.Fprintln(cmd.ErrOrStderr(), "Registered. You can now use 'lakebox ssh' to connect.")
+			blank(stderr)
+			fmt.Fprintf(stderr, "  Run %s to connect.\n\n", bold("lakebox ssh"))
 			return nil
 		},
 	}
diff --git a/cmd/lakebox/ssh.go b/cmd/lakebox/ssh.go
index 04a999bd40..483dbd38a8 100644
--- a/cmd/lakebox/ssh.go
+++ b/cmd/lakebox/ssh.go
@@ -53,7 +53,7 @@ Examples:
 			if _, err := os.Stat(keyPath); os.IsNotExist(err) {
 				return fmt.Errorf("lakebox SSH key not found at %s — run 'lakebox register' first", keyPath)
 			}
-			fmt.Fprintf(cmd.ErrOrStderr(), "Using SSH key: %s\n", keyPath)
+			stderr := cmd.ErrOrStderr()
 
 			// Parse args: everything before -- is the optional lakebox ID,
 			// everything after -- is passed through to ssh.
@@ -62,15 +62,12 @@ Examples:
 
 			dashAt := cmd.ArgsLenAtDash()
 			if dashAt == -1 {
-				// No -- found: first arg (if any) is lakebox ID.
 				if len(args) > 0 {
 					lakeboxID = args[0]
 				}
 			} else if dashAt == 0 {
-				// -- is first: no lakebox ID, rest is extra args.
 				extraArgs = args[dashAt:]
 			} else {
-				// lakebox ID before --, extra args after.
 				lakeboxID = args[0]
 				extraArgs = args[dashAt:]
 			}
@@ -79,7 +76,6 @@ Examples:
 			if lakeboxID == "" {
 				if def := getDefault(profile); def != "" {
 					lakeboxID = def
-					fmt.Fprintf(cmd.ErrOrStderr(), "Using default lakebox: %s\n", lakeboxID)
 				} else {
 					api := newLakeboxAPI(w)
 					pubKeyData, err := os.ReadFile(keyPath + ".pub")
@@ -87,22 +83,23 @@ Examples:
 						return fmt.Errorf("failed to read public key %s.pub: %w", keyPath, err)
 					}
 
-					fmt.Fprintf(cmd.ErrOrStderr(), "Creating lakebox...\n")
+					s := spin(stderr, "Provisioning your lakebox…")
 					result, err := api.create(ctx, string(pubKeyData))
 					if err != nil {
+						s.fail("Failed to create lakebox")
 						return fmt.Errorf("failed to create lakebox: %w", err)
 					}
 					lakeboxID = result.LakeboxID
-					fmt.Fprintf(cmd.ErrOrStderr(), "Lakebox %s created (status: %s)\n", lakeboxID, result.Status)
+					s.ok(fmt.Sprintf("Lakebox %s ready", bold(lakeboxID)))
 
 					if err := setDefault(profile, lakeboxID); err != nil {
-						fmt.Fprintf(cmd.ErrOrStderr(), "Warning: failed to save default: %v\n", err)
+						warn(stderr, fmt.Sprintf("Could not save default: %v", err))
 					}
 				}
 			}
 
-			fmt.Fprintf(cmd.ErrOrStderr(), "Connecting to %s@%s:%s...\n",
-				lakeboxID, gatewayHost, gatewayPort)
+			s := spin(stderr, fmt.Sprintf("Connecting to %s…", bold(lakeboxID)))
+			s.ok(fmt.Sprintf("Connected to %s", bold(lakeboxID)))
 			return execSSHDirect(lakeboxID, gatewayHost, gatewayPort, keyPath, extraArgs)
 		},
 	}
diff --git a/cmd/lakebox/status.go b/cmd/lakebox/status.go
index eaeeb8d7cc..bf2efbcaba 100644
--- a/cmd/lakebox/status.go
+++ b/cmd/lakebox/status.go
@@ -17,8 +17,8 @@ func newStatusCommand() *cobra.Command {
 		Long: `Show detailed status of a Lakebox environment.
 
 Example:
-  databricks lakebox status happy-panda-1234
-  databricks lakebox status happy-panda-1234 --json`,
+  lakebox status happy-panda-1234
+  lakebox status happy-panda-1234 --json`,
 		Args:    cobra.ExactArgs(1),
 		PreRunE: mustWorkspaceClient,
 		RunE: func(cmd *cobra.Command, args []string) error {
@@ -39,11 +39,14 @@ Example:
 				return enc.Encode(entry)
 			}
 
-			fmt.Fprintf(cmd.OutOrStdout(), "ID:     %s\n", extractLakeboxID(entry.Name))
-			fmt.Fprintf(cmd.OutOrStdout(), "Status: %s\n", entry.Status)
+			out := cmd.OutOrStdout()
+			blank(out)
+			field(out, "id", bold(extractLakeboxID(entry.Name)))
+			field(out, "status", status(entry.Status))
 			if entry.FQDN != "" {
-				fmt.Fprintf(cmd.OutOrStdout(), "FQDN:   %s\n", entry.FQDN)
+				field(out, "fqdn", dim(entry.FQDN))
 			}
+			blank(out)
 			return nil
 		},
 	}
diff --git a/cmd/lakebox/ui.go b/cmd/lakebox/ui.go
new file mode 100644
index 0000000000..2eab33310c
--- /dev/null
+++ b/cmd/lakebox/ui.go
@@ -0,0 +1,141 @@
+package lakebox
+
+import (
+	"fmt"
+	"io"
+	"os"
+	"strings"
+	"sync"
+	"time"
+)
+
+// Single accent color throughout. Bold for emphasis. Dim for metadata.
+const (
+	rs   = "\033[0m"  // reset
+	bo   = "\033[1m"  // bold
+	dm   = "\033[2m"  // dim
+	cyan = "\033[36m" // accent
+)
+
+func isTTY(w io.Writer) bool {
+	if f, ok := w.(*os.File); ok {
+		fi, err := f.Stat()
+		if err != nil {
+			return false
+		}
+		return fi.Mode()&os.ModeCharDevice != 0
+	}
+	return false
+}
+
+// spinner shows a braille spinner like Claude Code.
+type spinner struct {
+	w       io.Writer
+	msg     string
+	done    chan struct{}
+	once    sync.Once
+	started time.Time
+}
+
+func spin(w io.Writer, msg string) *spinner {
+	s := &spinner{w: w, msg: msg, done: make(chan struct{}), started: time.Now()}
+	if isTTY(w) {
+		go s.run()
+	} else {
+		fmt.Fprintf(w, "* %s\n", msg)
+	}
+	return s
+}
+
+func (s *spinner) run() {
+	frames := []string{"⣾", "⣽", "⣻", "⢿", "⡿", "⣟", "⣯", "⣷"}
+	i := 0
+	ticker := time.NewTicker(80 * time.Millisecond)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-s.done:
+			return
+		case <-ticker.C:
+			elapsed := time.Since(s.started).Truncate(time.Second)
+			fmt.Fprintf(s.w, "\r  %s%s%s %s%s%s %s(%s)%s  ",
+				cyan, frames[i%len(frames)], rs,
+				bo, s.msg, rs,
+				dm, elapsed, rs)
+			i++
+		}
+	}
+}
+
+func (s *spinner) ok(msg string) {
+	s.once.Do(func() {
+		close(s.done)
+		if isTTY(s.w) {
+			fmt.Fprintf(s.w, "\r\033[K  %s✓%s %s\n", cyan, rs, msg)
+		} else {
+			fmt.Fprintf(s.w, "✓ %s\n", msg)
+		}
+	})
+}
+
+func (s *spinner) fail(msg string) {
+	s.once.Do(func() {
+		close(s.done)
+		if isTTY(s.w) {
+			fmt.Fprintf(s.w, "\r\033[K  %s✗%s %s\n", cyan, rs, msg)
+		} else {
+			fmt.Fprintf(s.w, "✗ %s\n", msg)
+		}
+	})
+}
+
+// --- Consistent output primitives ---
+
+// status formats a status string with the accent color.
+func status(s string) string {
+	switch strings.ToLower(s) {
+	case "running":
+		return cyan + "running" + rs
+	case "stopped":
+		return dm + "stopped" + rs
+	case "creating":
+		return cyan + bo + "creating…" + rs
+	default:
+		return dm + strings.ToLower(s) + rs
+	}
+}
+
+// field prints "  label  value"
+func field(w io.Writer, label, value string) {
+	fmt.Fprintf(w, "  %s%-10s%s %s\n", dm, label, rs, value)
+}
+
+// ok prints "  ✓ message"
+func ok(w io.Writer, msg string) {
+	fmt.Fprintf(w, "  %s✓%s %s\n", cyan, rs, msg)
+}
+
+// warn prints "  ! message"
+func warn(w io.Writer, msg string) {
+	fmt.Fprintf(w, "  %s!%s %s\n", cyan, rs, msg)
+}
+
+// blank prints an empty line.
+func blank(w io.Writer) {
+	fmt.Fprintln(w)
+}
+
+// accent wraps text in the accent color.
+func accent(s string) string {
+	return cyan + s + rs
+}
+
+// bold wraps text in bold.
+func bold(s string) string {
+	return bo + s + rs
+}
+
+// dim wraps text in dim.
+func dim(s string) string {
+	return dm + s + rs
+}

From f9de7881c49d9351a3c839c274b0e763125f4006 Mon Sep 17 00:00:00 2001
From: shuochen0311 <shuo.chen@databricks.com>
Date: Wed, 29 Apr 2026 18:11:41 +0000
Subject: [PATCH 08/36] Fix CLI to match new lakebox API contract
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The lakebox manager moved its REST surface to a proto-defined service with
JSON transcoding (databricks-eng/universe#1839855 + follow-ups). That
changed three things this CLI was depending on:

1. JSON field name: each Lakebox message now serializes as `lakeboxId`
   (proto3 lowerCamelCase default), not `name`. List/status/create were
   parsing into `Name string \`json:"name"\`` and silently getting the
   empty string for every entry — the visible symptom was `lakebox list`
   showing rows with blank ID columns.

2. Status codes: proto-transcoded handlers return 200 OK uniformly. The
   CLI was checking 201 Created on POST /api/2.0/lakebox and 204
   NoContent on DELETE, both of which now look like errors.

3. Key registration moved to its own top-level collection at
   /api/2.0/lakebox-keys (was /api/2.0/lakebox/register-key), to avoid a
   path collision with /api/2.0/lakebox/{lakebox_id}.

Drop the now-unused `extractLakeboxID` helper — the wire field is the
customer-facing ID directly.

Verified against dev-aws-us-west-2: list, status, create, delete all
work end-to-end. register hits a separate manager-side issue (stale
UserKey records in TiDB that the new schema can't deserialize) — not
fixed here.

Co-authored-by: Isaac
---
 cmd/lakebox/api.go    | 36 +++++++++++++++++-------------------
 cmd/lakebox/list.go   |  4 ++--
 cmd/lakebox/status.go |  2 +-
 3 files changed, 20 insertions(+), 22 deletions(-)

diff --git a/cmd/lakebox/api.go b/cmd/lakebox/api.go
index 94877b4a42..04cbc1179c 100644
--- a/cmd/lakebox/api.go
+++ b/cmd/lakebox/api.go
@@ -25,16 +25,19 @@ type createRequest struct {
 }
 
 // createResponse is the JSON body returned by POST /api/2.0/lakebox.
+// Mirrors the `Lakebox` proto message after JSON transcoding.
 type createResponse struct {
-	LakeboxID string `json:"lakebox_id"`
+	LakeboxID string `json:"lakeboxId"`
 	Status    string `json:"status"`
+	FQDN      string `json:"fqdn"`
 }
 
 // lakeboxEntry is a single item in the list response.
+// Mirrors the `Lakebox` proto message after JSON transcoding.
 type lakeboxEntry struct {
-	Name   string `json:"name"`
-	Status string `json:"status"`
-	FQDN   string `json:"fqdn"`
+	LakeboxID string `json:"lakeboxId"`
+	Status    string `json:"status"`
+	FQDN      string `json:"fqdn"`
 }
 
 // listResponse is the JSON body returned by GET /api/2.0/lakebox.
@@ -70,7 +73,7 @@ func (a *lakeboxAPI) create(ctx context.Context, publicKey string) (*createRespo
 	}
 	defer resp.Body.Close()
 
-	if resp.StatusCode != http.StatusCreated {
+	if resp.StatusCode != http.StatusOK {
 		return nil, parseAPIError(resp)
 	}
 
@@ -127,7 +130,7 @@ func (a *lakeboxAPI) delete(ctx context.Context, id string) error {
 	}
 	defer resp.Body.Close()
 
-	if resp.StatusCode != http.StatusNoContent {
+	if resp.StatusCode != http.StatusOK {
 		return parseAPIError(resp)
 	}
 	return nil
@@ -163,12 +166,17 @@ func parseAPIError(resp *http.Response) error {
 	return fmt.Errorf("HTTP %d: %s", resp.StatusCode, string(body))
 }
 
-// registerKeyRequest is the JSON body for POST /api/2.0/lakebox/register-key.
+// User keys live at /api/2.0/lakebox-keys (separate top-level collection so
+// the path doesn't structurally overlap with /api/2.0/lakebox/{lakebox_id}).
+const lakeboxKeysAPIPath = "/api/2.0/lakebox-keys"
+
+// registerKeyRequest is the JSON body for POST /api/2.0/lakebox-keys.
 type registerKeyRequest struct {
 	PublicKey string `json:"public_key"`
+	Name      string `json:"name,omitempty"`
 }
 
-// registerKey calls POST /api/2.0/lakebox/register-key.
+// registerKey calls POST /api/2.0/lakebox-keys.
 func (a *lakeboxAPI) registerKey(ctx context.Context, publicKey string) error {
 	body := registerKeyRequest{PublicKey: publicKey}
 	jsonBody, err := json.Marshal(body)
@@ -176,7 +184,7 @@ func (a *lakeboxAPI) registerKey(ctx context.Context, publicKey string) error {
 		return fmt.Errorf("failed to marshal request: %w", err)
 	}
 
-	resp, err := a.doRequest(ctx, "POST", lakeboxAPIPath+"/register-key", bytes.NewReader(jsonBody))
+	resp, err := a.doRequest(ctx, "POST", lakeboxKeysAPIPath, bytes.NewReader(jsonBody))
 	if err != nil {
 		return err
 	}
@@ -187,13 +195,3 @@ func (a *lakeboxAPI) registerKey(ctx context.Context, publicKey string) error {
 	}
 	return nil
 }
-
-// extractLakeboxID extracts the short ID from a full resource name.
-// e.g. "apps/lakebox/instances/happy-panda-1234" -> "happy-panda-1234"
-func extractLakeboxID(name string) string {
-	parts := strings.Split(name, "/")
-	if len(parts) > 0 {
-		return parts[len(parts)-1]
-	}
-	return name
-}
diff --git a/cmd/lakebox/list.go b/cmd/lakebox/list.go
index 2ed3149658..69f9b2e3d7 100644
--- a/cmd/lakebox/list.go
+++ b/cmd/lakebox/list.go
@@ -56,7 +56,7 @@ Example:
 			// Compute column width.
 			col := 10
 			for _, e := range entries {
-				if l := len(extractLakeboxID(e.Name)); l > col {
+				if l := len(e.LakeboxID); l > col {
 					col = l
 				}
 			}
@@ -67,7 +67,7 @@ Example:
 			fmt.Fprintf(out, "  %s%s%s\n", dm, strings.Repeat("─", col+22), rs)
 
 			for _, e := range entries {
-				id := extractLakeboxID(e.Name)
+				id := e.LakeboxID
 				def := ""
 				if id == defaultID {
 					def = accent("*")
diff --git a/cmd/lakebox/status.go b/cmd/lakebox/status.go
index bf2efbcaba..d362143dc6 100644
--- a/cmd/lakebox/status.go
+++ b/cmd/lakebox/status.go
@@ -41,7 +41,7 @@ Example:
 
 			out := cmd.OutOrStdout()
 			blank(out)
-			field(out, "id", bold(extractLakeboxID(entry.Name)))
+			field(out, "id", bold(entry.LakeboxID))
 			field(out, "status", status(entry.Status))
 			if entry.FQDN != "" {
 				field(out, "fqdn", dim(entry.FQDN))

From 97e916e54a2ac8f1c9e454ae9616ee9bafac5f0e Mon Sep 17 00:00:00 2001
From: shuochen0311 <shuo.chen@databricks.com>
Date: Thu, 30 Apr 2026 17:04:39 +0000
Subject: [PATCH 09/36] Update CLI to lakebox sandbox/ssh-keys API surface
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Reynold's restructure (databricks-eng/universe#1874214) nested the two
lakebox resources under the service namespace — moving sandboxes from
/api/2.0/lakebox to /api/2.0/lakebox/sandboxes and SSH keys from
/api/2.0/lakebox-keys to /api/2.0/lakebox/ssh-keys — and renamed the
resource type from Lakebox to Sandbox, which surfaces on the wire as
sandboxId / sandboxes (was lakeboxId / lakeboxes).

CLI still pointed at the old paths and decoded the old field names, so
list / status / create returned empty IDs and 404s. Fix both endpoint
constants, rename the request/response types and fields to match the
proto, and update the four call sites in create / list / ssh / status.
User-facing copy ("Lakebox …") is unchanged — the product is still
Lakebox; only the resource type renamed.

Verified end-to-end against dev-aws-us-west-2: create / list / status
/ delete all work; ssh passthrough works.

Co-authored-by: Isaac
---
 cmd/lakebox/api.go    | 54 +++++++++++++++++++++++++------------------
 cmd/lakebox/create.go |  8 +++----
 cmd/lakebox/list.go   |  4 ++--
 cmd/lakebox/ssh.go    |  2 +-
 cmd/lakebox/status.go |  2 +-
 5 files changed, 39 insertions(+), 31 deletions(-)

diff --git a/cmd/lakebox/api.go b/cmd/lakebox/api.go
index 04cbc1179c..06b6de217b 100644
--- a/cmd/lakebox/api.go
+++ b/cmd/lakebox/api.go
@@ -12,37 +12,45 @@ import (
 	"github.com/databricks/databricks-sdk-go"
 )
 
-const lakeboxAPIPath = "/api/2.0/lakebox"
+// Sandboxes live under the `/sandboxes` sub-collection of the lakebox service
+// namespace (see `lakebox.proto` `LakeboxService.CreateSandbox`).
+const lakeboxAPIPath = "/api/2.0/lakebox/sandboxes"
 
 // lakeboxAPI wraps raw HTTP calls to the lakebox REST API.
 type lakeboxAPI struct {
 	w *databricks.WorkspaceClient
 }
 
-// createRequest is the JSON body for POST /api/2.0/lakebox.
+// createRequest is the JSON body for POST /api/2.0/lakebox/sandboxes.
+//
+// The proto-defined `CreateSandboxRequest` carries a `Sandbox sandbox = 1`
+// field today (every member is server-chosen), but JSON transcoding accepts
+// the unwrapped form for forward-compatible callers. Keep `public_key` here
+// as a no-op compat shim so older `lakebox create --public-key-file=...`
+// invocations don't error — the manager ignores it on the wire.
 type createRequest struct {
 	PublicKey string `json:"public_key,omitempty"`
 }
 
-// createResponse is the JSON body returned by POST /api/2.0/lakebox.
-// Mirrors the `Lakebox` proto message after JSON transcoding.
+// createResponse is the JSON body returned by POST /api/2.0/lakebox/sandboxes.
+// Mirrors the `Sandbox` proto message after JSON transcoding.
 type createResponse struct {
-	LakeboxID string `json:"lakeboxId"`
+	SandboxID string `json:"sandboxId"`
 	Status    string `json:"status"`
 	FQDN      string `json:"fqdn"`
 }
 
-// lakeboxEntry is a single item in the list response.
-// Mirrors the `Lakebox` proto message after JSON transcoding.
-type lakeboxEntry struct {
-	LakeboxID string `json:"lakeboxId"`
+// sandboxEntry is a single item in the list response.
+// Mirrors the `Sandbox` proto message after JSON transcoding.
+type sandboxEntry struct {
+	SandboxID string `json:"sandboxId"`
 	Status    string `json:"status"`
 	FQDN      string `json:"fqdn"`
 }
 
-// listResponse is the JSON body returned by GET /api/2.0/lakebox.
+// listResponse is the JSON body returned by GET /api/2.0/lakebox/sandboxes.
 type listResponse struct {
-	Lakeboxes []lakeboxEntry `json:"lakeboxes"`
+	Sandboxes []sandboxEntry `json:"sandboxes"`
 }
 
 // apiError is the error body returned by the lakebox API.
@@ -84,8 +92,8 @@ func (a *lakeboxAPI) create(ctx context.Context, publicKey string) (*createRespo
 	return &result, nil
 }
 
-// list calls GET /api/2.0/lakebox.
-func (a *lakeboxAPI) list(ctx context.Context) ([]lakeboxEntry, error) {
+// list calls GET /api/2.0/lakebox/sandboxes.
+func (a *lakeboxAPI) list(ctx context.Context) ([]sandboxEntry, error) {
 	resp, err := a.doRequest(ctx, "GET", lakeboxAPIPath, nil)
 	if err != nil {
 		return nil, err
@@ -100,11 +108,11 @@ func (a *lakeboxAPI) list(ctx context.Context) ([]lakeboxEntry, error) {
 	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
 		return nil, fmt.Errorf("failed to decode response: %w", err)
 	}
-	return result.Lakeboxes, nil
+	return result.Sandboxes, nil
 }
 
-// get calls GET /api/2.0/lakebox/{id}.
-func (a *lakeboxAPI) get(ctx context.Context, id string) (*lakeboxEntry, error) {
+// get calls GET /api/2.0/lakebox/sandboxes/{id}.
+func (a *lakeboxAPI) get(ctx context.Context, id string) (*sandboxEntry, error) {
 	resp, err := a.doRequest(ctx, "GET", lakeboxAPIPath+"/"+id, nil)
 	if err != nil {
 		return nil, err
@@ -115,14 +123,14 @@ func (a *lakeboxAPI) get(ctx context.Context, id string) (*lakeboxEntry, error)
 		return nil, parseAPIError(resp)
 	}
 
-	var result lakeboxEntry
+	var result sandboxEntry
 	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
 		return nil, fmt.Errorf("failed to decode response: %w", err)
 	}
 	return &result, nil
 }
 
-// delete calls DELETE /api/2.0/lakebox/{id}.
+// delete calls DELETE /api/2.0/lakebox/sandboxes/{id}.
 func (a *lakeboxAPI) delete(ctx context.Context, id string) error {
 	resp, err := a.doRequest(ctx, "DELETE", lakeboxAPIPath+"/"+id, nil)
 	if err != nil {
@@ -166,17 +174,17 @@ func parseAPIError(resp *http.Response) error {
 	return fmt.Errorf("HTTP %d: %s", resp.StatusCode, string(body))
 }
 
-// User keys live at /api/2.0/lakebox-keys (separate top-level collection so
-// the path doesn't structurally overlap with /api/2.0/lakebox/{lakebox_id}).
-const lakeboxKeysAPIPath = "/api/2.0/lakebox-keys"
+// SSH keys are now nested under the lakebox service namespace alongside
+// `sandboxes/` (see `LakeboxService.CreateSshKey`).
+const lakeboxKeysAPIPath = "/api/2.0/lakebox/ssh-keys"
 
-// registerKeyRequest is the JSON body for POST /api/2.0/lakebox-keys.
+// registerKeyRequest is the JSON body for POST /api/2.0/lakebox/ssh-keys.
 type registerKeyRequest struct {
 	PublicKey string `json:"public_key"`
 	Name      string `json:"name,omitempty"`
 }
 
-// registerKey calls POST /api/2.0/lakebox-keys.
+// registerKey calls POST /api/2.0/lakebox/ssh-keys.
 func (a *lakeboxAPI) registerKey(ctx context.Context, publicKey string) error {
 	body := registerKeyRequest{PublicKey: publicKey}
 	jsonBody, err := json.Marshal(body)
diff --git a/cmd/lakebox/create.go b/cmd/lakebox/create.go
index c4ce3a439e..096df26ce6 100644
--- a/cmd/lakebox/create.go
+++ b/cmd/lakebox/create.go
@@ -45,7 +45,7 @@ Example:
 				return fmt.Errorf("failed to create lakebox: %w", err)
 			}
 
-			s.ok(fmt.Sprintf("Lakebox %s is %s", bold(result.LakeboxID), status(result.Status)))
+			s.ok(fmt.Sprintf("Lakebox %s is %s", bold(result.SandboxID), status(result.Status)))
 
 			profile := w.Config.Profile
 			if profile == "" {
@@ -60,15 +60,15 @@ Example:
 				}
 			}
 			if shouldSetDefault {
-				if err := setDefault(profile, result.LakeboxID); err != nil {
+				if err := setDefault(profile, result.SandboxID); err != nil {
 					warn(stderr, fmt.Sprintf("Could not save default: %v", err))
 				} else {
-					field(stderr, "default", result.LakeboxID)
+					field(stderr, "default", result.SandboxID)
 				}
 			}
 
 			blank(stderr)
-			fmt.Fprintln(cmd.OutOrStdout(), result.LakeboxID)
+			fmt.Fprintln(cmd.OutOrStdout(), result.SandboxID)
 			return nil
 		},
 	}
diff --git a/cmd/lakebox/list.go b/cmd/lakebox/list.go
index 69f9b2e3d7..fe303028a0 100644
--- a/cmd/lakebox/list.go
+++ b/cmd/lakebox/list.go
@@ -56,7 +56,7 @@ Example:
 			// Compute column width.
 			col := 10
 			for _, e := range entries {
-				if l := len(e.LakeboxID); l > col {
+				if l := len(e.SandboxID); l > col {
 					col = l
 				}
 			}
@@ -67,7 +67,7 @@ Example:
 			fmt.Fprintf(out, "  %s%s%s\n", dm, strings.Repeat("─", col+22), rs)
 
 			for _, e := range entries {
-				id := e.LakeboxID
+				id := e.SandboxID
 				def := ""
 				if id == defaultID {
 					def = accent("*")
diff --git a/cmd/lakebox/ssh.go b/cmd/lakebox/ssh.go
index 483dbd38a8..11297f2786 100644
--- a/cmd/lakebox/ssh.go
+++ b/cmd/lakebox/ssh.go
@@ -89,7 +89,7 @@ Examples:
 						s.fail("Failed to create lakebox")
 						return fmt.Errorf("failed to create lakebox: %w", err)
 					}
-					lakeboxID = result.LakeboxID
+					lakeboxID = result.SandboxID
 					s.ok(fmt.Sprintf("Lakebox %s ready", bold(lakeboxID)))
 
 					if err := setDefault(profile, lakeboxID); err != nil {
diff --git a/cmd/lakebox/status.go b/cmd/lakebox/status.go
index d362143dc6..aa4a443d0a 100644
--- a/cmd/lakebox/status.go
+++ b/cmd/lakebox/status.go
@@ -41,7 +41,7 @@ Example:
 
 			out := cmd.OutOrStdout()
 			blank(out)
-			field(out, "id", bold(entry.LakeboxID))
+			field(out, "id", bold(entry.SandboxID))
 			field(out, "status", status(entry.Status))
 			if entry.FQDN != "" {
 				field(out, "fqdn", dim(entry.FQDN))

From 46642d1043589c5e48a3fffae1aaf224801b342d Mon Sep 17 00:00:00 2001
From: shuochen0311 <shuo.chen@databricks.com>
Date: Fri, 1 May 2026 03:36:49 +0000
Subject: [PATCH 10/36] Show auto-stop policy in lakebox list and status
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Surfaces the new per-sandbox auto-stop knobs the manager added
(databricks-eng/universe#1875183) so users can see at a glance how long
their sandbox will live before the watchdog reaps it.

- `sandboxEntry` gains pointer fields `IdleTimeoutSecs` and `Persist` so
  we keep the proto3 explicit-presence semantics ("not in response" vs
  "explicitly set to 0 / false").
- `autoStopLabel()` collapses the policy to one short token:
  - `persist == true` → `never`
  - `idle_timeout_secs > 0` → compact duration (`90s`, `15m`, `2h`,
    `1h30m`)
  - otherwise → the manager's global default (10m), rendered
    explicitly so the column never says `default`
- `lakebox list` adds an AUTOSTOP column between STATUS and DEFAULT.
- `lakebox status` adds an `autostop` field after `fqdn`.

Verified end-to-end against dev-aws-us-west-2 — list and status both
render `10m` for sandboxes with no per-record override.

Co-authored-by: Isaac
---
 cmd/lakebox/api.go    | 54 ++++++++++++++++++++++++++++++++++++++++---
 cmd/lakebox/list.go   | 21 +++++++++++++----
 cmd/lakebox/status.go |  1 +
 3 files changed, 69 insertions(+), 7 deletions(-)

diff --git a/cmd/lakebox/api.go b/cmd/lakebox/api.go
index 06b6de217b..3d0a4707b4 100644
--- a/cmd/lakebox/api.go
+++ b/cmd/lakebox/api.go
@@ -42,10 +42,58 @@ type createResponse struct {
 
 // sandboxEntry is a single item in the list response.
 // Mirrors the `Sandbox` proto message after JSON transcoding.
+//
+// IdleTimeoutSecs and Persist correspond to the proto's `optional` fields;
+// they're pointers so we can tell "field absent on the wire" (server has the
+// global default) from "explicitly set to 0 / false."
 type sandboxEntry struct {
-	SandboxID string `json:"sandboxId"`
-	Status    string `json:"status"`
-	FQDN      string `json:"fqdn"`
+	SandboxID       string `json:"sandboxId"`
+	Status          string `json:"status"`
+	FQDN            string `json:"fqdn"`
+	IdleTimeoutSecs *int64 `json:"idleTimeoutSecs,omitempty"`
+	Persist         *bool  `json:"persist,omitempty"`
+}
+
+// defaultAutoStopSecs mirrors the manager's `watchdog_idle_grace_secs`
+// fallback (10 minutes) used when a sandbox has no per-record override.
+// The value is also documented in `lakebox/CLAUDE.md` ("Sandbox
+// Watchdog" section). Hardcoded here so list/status can render the
+// effective timeout without an extra round-trip to fetch manager config.
+const defaultAutoStopSecs int64 = 600
+
+// autoStopLabel renders the auto-stop policy advertised by the manager
+// for one sandbox into a short human-readable string. Mirrors the wire
+// semantics from `lakebox/proto/lakebox.proto`:
+//   - `persist == true` → never auto-stops
+//   - `idle_timeout_secs` set and positive → that many seconds
+//   - otherwise → manager's global default (`defaultAutoStopSecs`)
+func (e *sandboxEntry) autoStopLabel() string {
+	if e.Persist != nil && *e.Persist {
+		return "never"
+	}
+	if e.IdleTimeoutSecs != nil && *e.IdleTimeoutSecs > 0 {
+		return formatDurationSecs(*e.IdleTimeoutSecs)
+	}
+	return formatDurationSecs(defaultAutoStopSecs)
+}
+
+// formatDurationSecs prints `secs` as a compact duration (e.g. `90s`,
+// `15m`, `2h`, `1h30m`). Falls back to seconds if it's not a clean
+// minute/hour multiple. Avoids pulling in a dependency just for this.
+func formatDurationSecs(secs int64) string {
+	if secs < 60 {
+		return fmt.Sprintf("%ds", secs)
+	}
+	if secs%3600 == 0 {
+		return fmt.Sprintf("%dh", secs/3600)
+	}
+	if secs >= 3600 {
+		return fmt.Sprintf("%dh%dm", secs/3600, (secs%3600)/60)
+	}
+	if secs%60 == 0 {
+		return fmt.Sprintf("%dm", secs/60)
+	}
+	return fmt.Sprintf("%ds", secs)
 }
 
 // listResponse is the JSON body returned by GET /api/2.0/lakebox/sandboxes.
diff --git a/cmd/lakebox/list.go b/cmd/lakebox/list.go
index fe303028a0..f058524e7e 100644
--- a/cmd/lakebox/list.go
+++ b/cmd/lakebox/list.go
@@ -53,18 +53,25 @@ Example:
 
 			out := cmd.OutOrStdout()
 
-			// Compute column width.
+			// Compute column widths. AUTOSTOP holds short tokens like
+			// `default`, `never`, `15m`, `1h30m` — 8 chars covers them.
 			col := 10
+			autostopCol := 8
 			for _, e := range entries {
 				if l := len(e.SandboxID); l > col {
 					col = l
 				}
+				if l := len(e.autoStopLabel()); l > autostopCol {
+					autostopCol = l
+				}
 			}
 			col += 2
+			autostopCol += 2
 
 			blank(out)
-			fmt.Fprintf(out, "  %s%-*s  %-10s  %s%s\n", dm, col, "ID", "STATUS", "DEFAULT", rs)
-			fmt.Fprintf(out, "  %s%s%s\n", dm, strings.Repeat("─", col+22), rs)
+			fmt.Fprintf(out, "  %s%-*s  %-10s  %-*s  %s%s\n",
+				dm, col, "ID", "STATUS", autostopCol, "AUTOSTOP", "DEFAULT", rs)
+			fmt.Fprintf(out, "  %s%s%s\n", dm, strings.Repeat("─", col+10+autostopCol+12), rs)
 
 			for _, e := range entries {
 				id := e.SandboxID
@@ -83,13 +90,19 @@ Example:
 				if stPad < 0 {
 					stPad = 0
 				}
+				as := e.autoStopLabel()
+				asPad := autostopCol - len(as)
+				if asPad < 0 {
+					asPad = 0
+				}
 				idStr := bold(id)
 				if strings.EqualFold(e.Status, "running") {
 					idStr = cyan + bo + id + rs
 				}
-				fmt.Fprintf(out, "  %s%s  %s%s  %s\n",
+				fmt.Fprintf(out, "  %s%s  %s%s  %s%s  %s\n",
 					idStr, strings.Repeat(" ", idPad),
 					st, strings.Repeat(" ", stPad),
+					dim(as), strings.Repeat(" ", asPad),
 					def)
 			}
 			blank(out)
diff --git a/cmd/lakebox/status.go b/cmd/lakebox/status.go
index aa4a443d0a..f5df1ee4a4 100644
--- a/cmd/lakebox/status.go
+++ b/cmd/lakebox/status.go
@@ -46,6 +46,7 @@ Example:
 			if entry.FQDN != "" {
 				field(out, "fqdn", dim(entry.FQDN))
 			}
+			field(out, "autostop", dim(entry.autoStopLabel()))
 			blank(out)
 			return nil
 		},

From 412ff70988517e1e090ba539fa4d295b8ea7de43 Mon Sep 17 00:00:00 2001
From: shuochen0311 <shuo.chen@databricks.com>
Date: Fri, 1 May 2026 04:17:27 +0000
Subject: [PATCH 11/36] Add lakebox config command for setting auto-stop policy
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Surfaces the per-sandbox auto-stop knobs the manager added in
databricks-eng/universe#1875183 so users can flip them from the CLI
instead of curl + JSON.

  lakebox config <id> --idle-timeout 15m       # 15-minute timeout
  lakebox config <id> --idle-timeout 1h30m     # any Go duration
  lakebox config <id> --idle-timeout 0         # clear → manager default
  lakebox config <id> --persist                # never auto-stop
  lakebox config <id> --persist=false          # back to timeout path
  lakebox config <id> --idle-timeout 30m --persist=false   # combined

Implementation notes:

- `updateBody` is the inner Sandbox sent in the PATCH body. The proto's
  `(google.api.http)` declares `body: "sandbox"`, so the HTTP body is
  the inner `Sandbox` message, NOT a `{"sandbox": {...}}` envelope.
  First wired-up version got this wrong and the manager rejected with
  "unknown field `sandbox`" — kept the type comment to flag the gotcha
  for the next reader.
- `IdleTimeoutSecs` carries `,string` JSON tag because proto3 JSON
  canonical form serializes int64 as a quoted string. The manager
  accepts both bare-number and quoted-string on input but always
  emits quoted on output, so without the tag we hit "cannot unmarshal
  string into Go struct field … int64" on the response read-back.
- Pointer fields (`*int64`, `*bool`) carry proto3 explicit-presence
  through to the wire — only the flags the user actually passed get
  emitted, so a `--persist`-only invocation does not clobber an
  existing idle_timeout (and vice-versa).
- Client-side range pre-flight (`[60s, 86400s]` plus the 0 clear
  sentinel) mirrors the manager's `MIN_IDLE_TIMEOUT_SECS` /
  `MAX_IDLE_TIMEOUT_SECS` constants so users get a clearer error
  than the server's `INVALID_ARGUMENT`.

Verified end-to-end against dev-aws-us-west-2:
  config --idle-timeout 15m  → status shows `15m`
  config --persist           → status shows `never`
  config --idle-timeout 0 --persist=false → status shows `10m`

Co-authored-by: Isaac
---
 cmd/lakebox/api.go     |  60 ++++++++++++++++++-
 cmd/lakebox/config.go  | 128 +++++++++++++++++++++++++++++++++++++++++
 cmd/lakebox/lakebox.go |   1 +
 3 files changed, 188 insertions(+), 1 deletion(-)
 create mode 100644 cmd/lakebox/config.go

diff --git a/cmd/lakebox/api.go b/cmd/lakebox/api.go
index 3d0a4707b4..73d4b5d81f 100644
--- a/cmd/lakebox/api.go
+++ b/cmd/lakebox/api.go
@@ -46,11 +46,15 @@ type createResponse struct {
 // IdleTimeoutSecs and Persist correspond to the proto's `optional` fields;
 // they're pointers so we can tell "field absent on the wire" (server has the
 // global default) from "explicitly set to 0 / false."
+//
+// `IdleTimeoutSecs` carries a `,string` JSON tag because proto3 JSON
+// canonical form serializes int64 as a quoted string. The field is read
+// off the wire as `"900"`, not `900`.
 type sandboxEntry struct {
 	SandboxID       string `json:"sandboxId"`
 	Status          string `json:"status"`
 	FQDN            string `json:"fqdn"`
-	IdleTimeoutSecs *int64 `json:"idleTimeoutSecs,omitempty"`
+	IdleTimeoutSecs *int64 `json:"idleTimeoutSecs,omitempty,string"`
 	Persist         *bool  `json:"persist,omitempty"`
 }
 
@@ -178,6 +182,60 @@ func (a *lakeboxAPI) get(ctx context.Context, id string) (*sandboxEntry, error)
 	return &result, nil
 }
 
+// updateBody is the PATCH request body. The proto declares
+// `UpdateSandboxRequest { Sandbox sandbox = 1 }` with `body: "sandbox"`
+// in the (google.api.http) annotation, so the HTTP body is the inner
+// `Sandbox` message directly — there is no `{"sandbox": {...}}`
+// wrapping on the wire.
+//
+// Pointer fields encode the proto3 `optional` semantics — only the
+// fields we explicitly set are emitted, leaving everything else
+// server-untouched.
+type updateBody struct {
+	SandboxID string `json:"sandbox_id"`
+	// `,string` matches proto3 JSON canonical encoding; the manager
+	// accepts both quoted-string and bare-number int64 on input.
+	IdleTimeoutSecs *int64 `json:"idle_timeout_secs,omitempty,string"`
+	Persist         *bool  `json:"persist,omitempty"`
+}
+
+// update calls PATCH /api/2.0/lakebox/sandboxes/{id} with whichever of
+// `idle_timeout_secs` / `persist` the caller chose to set. Fields left
+// nil are omitted from the wire payload, so the server preserves their
+// current values. Returns the refreshed `sandboxEntry`.
+func (a *lakeboxAPI) update(
+	ctx context.Context,
+	id string,
+	idleTimeoutSecs *int64,
+	persist *bool,
+) (*sandboxEntry, error) {
+	body := updateBody{
+		SandboxID:       id,
+		IdleTimeoutSecs: idleTimeoutSecs,
+		Persist:         persist,
+	}
+	jsonBody, err := json.Marshal(body)
+	if err != nil {
+		return nil, fmt.Errorf("failed to marshal request: %w", err)
+	}
+
+	resp, err := a.doRequest(ctx, "PATCH", lakeboxAPIPath+"/"+id, bytes.NewReader(jsonBody))
+	if err != nil {
+		return nil, err
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return nil, parseAPIError(resp)
+	}
+
+	var result sandboxEntry
+	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
+		return nil, fmt.Errorf("failed to decode response: %w", err)
+	}
+	return &result, nil
+}
+
 // delete calls DELETE /api/2.0/lakebox/sandboxes/{id}.
 func (a *lakeboxAPI) delete(ctx context.Context, id string) error {
 	resp, err := a.doRequest(ctx, "DELETE", lakeboxAPIPath+"/"+id, nil)
diff --git a/cmd/lakebox/config.go b/cmd/lakebox/config.go
new file mode 100644
index 0000000000..e909489f2d
--- /dev/null
+++ b/cmd/lakebox/config.go
@@ -0,0 +1,128 @@
+package lakebox
+
+import (
+	"fmt"
+	"time"
+
+	"github.com/databricks/cli/libs/cmdctx"
+	"github.com/spf13/cobra"
+)
+
+// MIN_IDLE_TIMEOUT_SECS / MAX_IDLE_TIMEOUT_SECS mirror the manager-side
+// constants in lakebox/src/api/handlers/sandbox.rs. Pre-flighting client-side
+// gives a clearer error than waiting for the server's INVALID_ARGUMENT.
+const (
+	minIdleTimeoutSecs = 60
+	maxIdleTimeoutSecs = 86_400
+)
+
+func newConfigCommand() *cobra.Command {
+	var idleTimeoutFlag string
+	var persistFlag bool
+
+	cmd := &cobra.Command{
+		Use:   "config <lakebox-id>",
+		Short: "Configure a Lakebox's auto-stop policy",
+		Long: `Configure a Lakebox's auto-stop policy.
+
+Two knobs are independent — pass either or both:
+
+  --idle-timeout <duration>   Per-sandbox idle timeout. The watchdog reaps
+                              the sandbox after this much idle time. Pass
+                              0 (or 0s) to clear and revert to the manager's
+                              global default (10m). Valid range when set:
+                              60s to 24h.
+
+  --persist[=true|false]      When true, the sandbox is exempt from
+                              idle-driven auto-stop entirely. The
+                              --idle-timeout setting is ignored while
+                              persist is on. Sandbox still stops on
+                              explicit 'lakebox delete'.
+
+Examples:
+  lakebox config happy-panda-1234 --idle-timeout 15m
+  lakebox config happy-panda-1234 --idle-timeout 1h30m
+  lakebox config happy-panda-1234 --idle-timeout 0           # clear, use default
+  lakebox config happy-panda-1234 --persist                  # never auto-stop
+  lakebox config happy-panda-1234 --persist=false            # back to timeout path
+  lakebox config happy-panda-1234 --idle-timeout 30m --persist=false`,
+		Args:    cobra.ExactArgs(1),
+		PreRunE: mustWorkspaceClient,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			ctx := cmd.Context()
+			w := cmdctx.WorkspaceClient(ctx)
+			api := newLakeboxAPI(w)
+			out := cmd.OutOrStdout()
+
+			id := args[0]
+
+			// Translate flag presence + value into the proto3
+			// optional-field semantics the server expects.
+			var idleSecs *int64
+			if cmd.Flags().Changed("idle-timeout") {
+				secs, err := parseIdleTimeoutFlag(idleTimeoutFlag)
+				if err != nil {
+					return err
+				}
+				idleSecs = &secs
+			}
+
+			var persist *bool
+			if cmd.Flags().Changed("persist") {
+				p := persistFlag
+				persist = &p
+			}
+
+			if idleSecs == nil && persist == nil {
+				return fmt.Errorf("nothing to update — pass --idle-timeout and/or --persist")
+			}
+
+			updated, err := api.update(ctx, id, idleSecs, persist)
+			if err != nil {
+				return fmt.Errorf("failed to update lakebox %s: %w", id, err)
+			}
+
+			blank(out)
+			field(out, "id", bold(updated.SandboxID))
+			field(out, "autostop", dim(updated.autoStopLabel()))
+			blank(out)
+			return nil
+		},
+	}
+
+	cmd.Flags().StringVar(&idleTimeoutFlag, "idle-timeout", "",
+		"Idle timeout (e.g. 15m, 1h30m, 90s). Pass 0 to clear and revert to the manager's default.")
+	cmd.Flags().BoolVar(&persistFlag, "persist", false,
+		"When true, this sandbox never auto-stops on idle. Pass --persist=false to revert.")
+
+	return cmd
+}
+
+// parseIdleTimeoutFlag accepts the same syntax as time.ParseDuration plus
+// the special-case "0" / "0s" → clear. Anything else outside the
+// [60s, 86400s] window is rejected client-side.
+func parseIdleTimeoutFlag(raw string) (int64, error) {
+	d, err := time.ParseDuration(raw)
+	if err != nil {
+		// Allow bare integer seconds as a convenience (`--idle-timeout 900`).
+		var secs int64
+		if _, e2 := fmt.Sscanf(raw, "%d", &secs); e2 == nil {
+			return checkIdleSecs(secs)
+		}
+		return 0, fmt.Errorf("invalid --idle-timeout %q: %w (use Go duration syntax, e.g. 15m, 1h30m)", raw, err)
+	}
+	return checkIdleSecs(int64(d.Seconds()))
+}
+
+func checkIdleSecs(secs int64) (int64, error) {
+	if secs == 0 {
+		return 0, nil // clear / revert to global default
+	}
+	if secs < minIdleTimeoutSecs || secs > maxIdleTimeoutSecs {
+		return 0, fmt.Errorf(
+			"idle-timeout must be 0 (clear) or between %ds and %ds, got %ds",
+			minIdleTimeoutSecs, maxIdleTimeoutSecs, secs,
+		)
+	}
+	return secs, nil
+}
diff --git a/cmd/lakebox/lakebox.go b/cmd/lakebox/lakebox.go
index 6a968df87a..25a9b479e5 100644
--- a/cmd/lakebox/lakebox.go
+++ b/cmd/lakebox/lakebox.go
@@ -39,6 +39,7 @@ The CLI manages your ~/.ssh/config so you can also connect directly:
 	cmd.AddCommand(newCreateCommand())
 	cmd.AddCommand(newDeleteCommand())
 	cmd.AddCommand(newStatusCommand())
+	cmd.AddCommand(newConfigCommand())
 
 	return cmd
 }

From 03a6240531e637a22cf05f97c166cf18e5dee19a Mon Sep 17 00:00:00 2001
From: shuochen0311 <shuo.chen@databricks.com>
Date: Fri, 1 May 2026 04:34:16 +0000
Subject: [PATCH 12/36] =?UTF-8?q?Rename=20persist=20=E2=86=92=20no=5Fautos?=
 =?UTF-8?q?top=20and=20document=20auto-clear=20behavior?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Tracks the matching rename in the lakebox manager
(databricks-eng/universe#1875183 follow-up). The manager-side flag
moved from `persist` to `no_autostop` because the original name
conflicted with the storage-persistence concept already in this
codebase.

CLI changes:
  --persist          → --no-autostop
  --persist=false    → --no-autostop=false

Plus a help-text note on the manager's new auto-clear behavior:
setting `--idle-timeout` to a non-zero value in a follow-up call
clears `--no-autostop` automatically, on the assumption that the
caller wants timeout-based stopping back. The CLI itself does not
need any extra logic for this — the manager handles it server-side
based on field presence in the PATCH body, and the CLI's existing
"omit unset flags from the wire payload" semantics (proto3
explicit-presence via *bool / *int64) feed straight into that.

Verified the marshal output matches what the new manager expects:
  --no-autostop      → {"sandbox_id":"x","no_autostop":true}
  --idle-timeout 15m → {"sandbox_id":"x","idle_timeout_secs":"900"}
  no flags           → {"sandbox_id":"x"}                      (rejected)

End-to-end against staging blocked until the manager PR rolls out.

Co-authored-by: Isaac
---
 cmd/lakebox/api.go    | 16 ++++++++--------
 cmd/lakebox/config.go | 34 ++++++++++++++++++----------------
 2 files changed, 26 insertions(+), 24 deletions(-)

diff --git a/cmd/lakebox/api.go b/cmd/lakebox/api.go
index 73d4b5d81f..288e704a9e 100644
--- a/cmd/lakebox/api.go
+++ b/cmd/lakebox/api.go
@@ -43,7 +43,7 @@ type createResponse struct {
 // sandboxEntry is a single item in the list response.
 // Mirrors the `Sandbox` proto message after JSON transcoding.
 //
-// IdleTimeoutSecs and Persist correspond to the proto's `optional` fields;
+// IdleTimeoutSecs and NoAutostop correspond to the proto's `optional` fields;
 // they're pointers so we can tell "field absent on the wire" (server has the
 // global default) from "explicitly set to 0 / false."
 //
@@ -55,7 +55,7 @@ type sandboxEntry struct {
 	Status          string `json:"status"`
 	FQDN            string `json:"fqdn"`
 	IdleTimeoutSecs *int64 `json:"idleTimeoutSecs,omitempty,string"`
-	Persist         *bool  `json:"persist,omitempty"`
+	NoAutostop      *bool  `json:"noAutostop,omitempty"`
 }
 
 // defaultAutoStopSecs mirrors the manager's `watchdog_idle_grace_secs`
@@ -68,11 +68,11 @@ const defaultAutoStopSecs int64 = 600
 // autoStopLabel renders the auto-stop policy advertised by the manager
 // for one sandbox into a short human-readable string. Mirrors the wire
 // semantics from `lakebox/proto/lakebox.proto`:
-//   - `persist == true` → never auto-stops
+//   - `no_autostop == true` → never auto-stops
 //   - `idle_timeout_secs` set and positive → that many seconds
 //   - otherwise → manager's global default (`defaultAutoStopSecs`)
 func (e *sandboxEntry) autoStopLabel() string {
-	if e.Persist != nil && *e.Persist {
+	if e.NoAutostop != nil && *e.NoAutostop {
 		return "never"
 	}
 	if e.IdleTimeoutSecs != nil && *e.IdleTimeoutSecs > 0 {
@@ -196,23 +196,23 @@ type updateBody struct {
 	// `,string` matches proto3 JSON canonical encoding; the manager
 	// accepts both quoted-string and bare-number int64 on input.
 	IdleTimeoutSecs *int64 `json:"idle_timeout_secs,omitempty,string"`
-	Persist         *bool  `json:"persist,omitempty"`
+	NoAutostop      *bool  `json:"no_autostop,omitempty"`
 }
 
 // update calls PATCH /api/2.0/lakebox/sandboxes/{id} with whichever of
-// `idle_timeout_secs` / `persist` the caller chose to set. Fields left
+// `idle_timeout_secs` / `no_autostop` the caller chose to set. Fields left
 // nil are omitted from the wire payload, so the server preserves their
 // current values. Returns the refreshed `sandboxEntry`.
 func (a *lakeboxAPI) update(
 	ctx context.Context,
 	id string,
 	idleTimeoutSecs *int64,
-	persist *bool,
+	noAutostop *bool,
 ) (*sandboxEntry, error) {
 	body := updateBody{
 		SandboxID:       id,
 		IdleTimeoutSecs: idleTimeoutSecs,
-		Persist:         persist,
+		NoAutostop:      noAutostop,
 	}
 	jsonBody, err := json.Marshal(body)
 	if err != nil {
diff --git a/cmd/lakebox/config.go b/cmd/lakebox/config.go
index e909489f2d..fe3b80ddf2 100644
--- a/cmd/lakebox/config.go
+++ b/cmd/lakebox/config.go
@@ -18,7 +18,7 @@ const (
 
 func newConfigCommand() *cobra.Command {
 	var idleTimeoutFlag string
-	var persistFlag bool
+	var noAutostopFlag bool
 
 	cmd := &cobra.Command{
 		Use:   "config <lakebox-id>",
@@ -33,19 +33,21 @@ Two knobs are independent — pass either or both:
                               global default (10m). Valid range when set:
                               60s to 24h.
 
-  --persist[=true|false]      When true, the sandbox is exempt from
+  --no-autostop[=true|false]  When true, the sandbox is exempt from
                               idle-driven auto-stop entirely. The
                               --idle-timeout setting is ignored while
-                              persist is on. Sandbox still stops on
-                              explicit 'lakebox delete'.
+                              this is on. Setting --idle-timeout to a
+                              non-zero value in a later call clears
+                              --no-autostop automatically. Sandbox still
+                              stops on explicit 'lakebox delete'.
 
 Examples:
   lakebox config happy-panda-1234 --idle-timeout 15m
   lakebox config happy-panda-1234 --idle-timeout 1h30m
   lakebox config happy-panda-1234 --idle-timeout 0           # clear, use default
-  lakebox config happy-panda-1234 --persist                  # never auto-stop
-  lakebox config happy-panda-1234 --persist=false            # back to timeout path
-  lakebox config happy-panda-1234 --idle-timeout 30m --persist=false`,
+  lakebox config happy-panda-1234 --no-autostop                  # never auto-stop
+  lakebox config happy-panda-1234 --no-autostop=false            # back to timeout path
+  lakebox config happy-panda-1234 --idle-timeout 30m --no-autostop=false`,
 		Args:    cobra.ExactArgs(1),
 		PreRunE: mustWorkspaceClient,
 		RunE: func(cmd *cobra.Command, args []string) error {
@@ -67,17 +69,17 @@ Examples:
 				idleSecs = &secs
 			}
 
-			var persist *bool
-			if cmd.Flags().Changed("persist") {
-				p := persistFlag
-				persist = &p
+			var noAutostop *bool
+			if cmd.Flags().Changed("no-autostop") {
+				p := noAutostopFlag
+				noAutostop = &p
 			}
 
-			if idleSecs == nil && persist == nil {
-				return fmt.Errorf("nothing to update — pass --idle-timeout and/or --persist")
+			if idleSecs == nil && noAutostop == nil {
+				return fmt.Errorf("nothing to update — pass --idle-timeout and/or --no-autostop")
 			}
 
-			updated, err := api.update(ctx, id, idleSecs, persist)
+			updated, err := api.update(ctx, id, idleSecs, noAutostop)
 			if err != nil {
 				return fmt.Errorf("failed to update lakebox %s: %w", id, err)
 			}
@@ -92,8 +94,8 @@ Examples:
 
 	cmd.Flags().StringVar(&idleTimeoutFlag, "idle-timeout", "",
 		"Idle timeout (e.g. 15m, 1h30m, 90s). Pass 0 to clear and revert to the manager's default.")
-	cmd.Flags().BoolVar(&persistFlag, "persist", false,
-		"When true, this sandbox never auto-stops on idle. Pass --persist=false to revert.")
+	cmd.Flags().BoolVar(&noAutostopFlag, "no-autostop", false,
+		"When true, this sandbox never auto-stops on idle. Pass --no-autostop=false to revert.")
 
 	return cmd
 }

From 8cfe3bb939cfb7639fc1da28e96487f5d2a37f9b Mon Sep 17 00:00:00 2001
From: shuochen0311 <shuo.chen@databricks.com>
Date: Fri, 1 May 2026 05:15:53 +0000
Subject: [PATCH 13/36] Switch idle_timeout wire type to
 google.protobuf.Duration
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Tracks the matching change in the lakebox manager
(databricks-eng/universe#1875183) which moved the per-sandbox idle
timeout off `optional int64 idle_timeout_secs = 7` and onto
`optional google.protobuf.Duration idle_timeout = 7`. Drops the
sentinel-overloaded int64 in favor of a duration-typed field.

Wire shape:
- Response field is now `idleTimeout` carrying a proto3-canonical
  Duration string (e.g. `"900s"`); parsed into seconds via
  `time.ParseDuration` for the autostop column.
- Request body sends `idle_timeout` as the same string format.

The CLI flag stays `--idle-timeout` (Go duration string in / Go
duration string out); only the wire encoding changes. `list` and
`status` show the manager's global default for any sandbox whose
per-record value isn't yet visible under the new field name — that's
deliberate forward-compat behavior so an older manager + newer CLI
combination just degrades to showing the default rather than crashing.

Co-authored-by: Isaac
---
 cmd/lakebox/api.go | 73 +++++++++++++++++++++++++++++++---------------
 1 file changed, 49 insertions(+), 24 deletions(-)

diff --git a/cmd/lakebox/api.go b/cmd/lakebox/api.go
index 288e704a9e..4fb3af2630 100644
--- a/cmd/lakebox/api.go
+++ b/cmd/lakebox/api.go
@@ -8,6 +8,7 @@ import (
 	"io"
 	"net/http"
 	"strings"
+	"time"
 
 	"github.com/databricks/databricks-sdk-go"
 )
@@ -43,19 +44,38 @@ type createResponse struct {
 // sandboxEntry is a single item in the list response.
 // Mirrors the `Sandbox` proto message after JSON transcoding.
 //
-// IdleTimeoutSecs and NoAutostop correspond to the proto's `optional` fields;
-// they're pointers so we can tell "field absent on the wire" (server has the
-// global default) from "explicitly set to 0 / false."
+// IdleTimeout and NoAutostop correspond to the proto's `optional` fields;
+// they're pointers so we can tell "field absent on the wire" (server has
+// the global default) from "explicitly set to 0 / false."
 //
-// `IdleTimeoutSecs` carries a `,string` JSON tag because proto3 JSON
-// canonical form serializes int64 as a quoted string. The field is read
-// off the wire as `"900"`, not `900`.
+// `IdleTimeout` is a `google.protobuf.Duration`. Proto3 JSON canonical
+// form serializes Duration as a string with an `s` suffix (e.g.
+// `"900s"`), so the Go field is `*string` and we parse on read.
 type sandboxEntry struct {
-	SandboxID       string `json:"sandboxId"`
-	Status          string `json:"status"`
-	FQDN            string `json:"fqdn"`
-	IdleTimeoutSecs *int64 `json:"idleTimeoutSecs,omitempty,string"`
-	NoAutostop      *bool  `json:"noAutostop,omitempty"`
+	SandboxID   string  `json:"sandboxId"`
+	Status      string  `json:"status"`
+	FQDN        string  `json:"fqdn"`
+	IdleTimeout *string `json:"idleTimeout,omitempty"`
+	NoAutostop  *bool   `json:"noAutostop,omitempty"`
+}
+
+// idleTimeoutSecs parses the proto3-canonical Duration string off
+// `IdleTimeout` (e.g. `"900s"` → `900`). Returns 0 when unset or when
+// the string is not a recognizable Duration. Sub-second precision is
+// dropped — the watchdog only acts on whole seconds.
+func (e *sandboxEntry) idleTimeoutSecs() int64 {
+	if e.IdleTimeout == nil {
+		return 0
+	}
+	s := *e.IdleTimeout
+	if !strings.HasSuffix(s, "s") {
+		return 0
+	}
+	d, err := time.ParseDuration(s)
+	if err != nil {
+		return 0
+	}
+	return int64(d.Seconds())
 }
 
 // defaultAutoStopSecs mirrors the manager's `watchdog_idle_grace_secs`
@@ -69,14 +89,14 @@ const defaultAutoStopSecs int64 = 600
 // for one sandbox into a short human-readable string. Mirrors the wire
 // semantics from `lakebox/proto/lakebox.proto`:
 //   - `no_autostop == true` → never auto-stops
-//   - `idle_timeout_secs` set and positive → that many seconds
+//   - `idle_timeout` set and positive → that many seconds
 //   - otherwise → manager's global default (`defaultAutoStopSecs`)
 func (e *sandboxEntry) autoStopLabel() string {
 	if e.NoAutostop != nil && *e.NoAutostop {
 		return "never"
 	}
-	if e.IdleTimeoutSecs != nil && *e.IdleTimeoutSecs > 0 {
-		return formatDurationSecs(*e.IdleTimeoutSecs)
+	if secs := e.idleTimeoutSecs(); secs > 0 {
+		return formatDurationSecs(secs)
 	}
 	return formatDurationSecs(defaultAutoStopSecs)
 }
@@ -190,17 +210,17 @@ func (a *lakeboxAPI) get(ctx context.Context, id string) (*sandboxEntry, error)
 //
 // Pointer fields encode the proto3 `optional` semantics — only the
 // fields we explicitly set are emitted, leaving everything else
-// server-untouched.
+// server-untouched. `IdleTimeout` is a proto3-canonical Duration
+// string (e.g. `"900s"`); the server-side wire type is
+// `google.protobuf.Duration`.
 type updateBody struct {
-	SandboxID string `json:"sandbox_id"`
-	// `,string` matches proto3 JSON canonical encoding; the manager
-	// accepts both quoted-string and bare-number int64 on input.
-	IdleTimeoutSecs *int64 `json:"idle_timeout_secs,omitempty,string"`
-	NoAutostop      *bool  `json:"no_autostop,omitempty"`
+	SandboxID   string  `json:"sandbox_id"`
+	IdleTimeout *string `json:"idle_timeout,omitempty"`
+	NoAutostop  *bool   `json:"no_autostop,omitempty"`
 }
 
 // update calls PATCH /api/2.0/lakebox/sandboxes/{id} with whichever of
-// `idle_timeout_secs` / `no_autostop` the caller chose to set. Fields left
+// `idle_timeout` / `no_autostop` the caller chose to set. Fields left
 // nil are omitted from the wire payload, so the server preserves their
 // current values. Returns the refreshed `sandboxEntry`.
 func (a *lakeboxAPI) update(
@@ -209,10 +229,15 @@ func (a *lakeboxAPI) update(
 	idleTimeoutSecs *int64,
 	noAutostop *bool,
 ) (*sandboxEntry, error) {
+	var idleTimeout *string
+	if idleTimeoutSecs != nil {
+		s := fmt.Sprintf("%ds", *idleTimeoutSecs)
+		idleTimeout = &s
+	}
 	body := updateBody{
-		SandboxID:       id,
-		IdleTimeoutSecs: idleTimeoutSecs,
-		NoAutostop:      noAutostop,
+		SandboxID:   id,
+		IdleTimeout: idleTimeout,
+		NoAutostop:  noAutostop,
 	}
 	jsonBody, err := json.Marshal(body)
 	if err != nil {

From b87b71291900aacf876b0661914958450745bbe9 Mon Sep 17 00:00:00 2001
From: shuochen0311 <shuo.chen@databricks.com>
Date: Sat, 2 May 2026 04:36:17 +0000
Subject: [PATCH 14/36] [lakebox] Support staging workspaces in CLI ssh + api
 routing

- ssh: auto-pick uw2.s.dbrx.dev when the workspace host has `.staging.` in it,
  otherwise keep using prod uw2.dbrx.dev. `--gateway` still overrides.
- api: when the workspace host carries a `?o=<id>` selector or the SDK config
  has a workspace_id, send `X-Databricks-Org-Id` so multi-workspace gateways
  (dogfood.staging.databricks.com) route the request to the right workspace.
  Without it the gateway rejects PATs with "Credential was not sent or was
  of an unsupported type for this API".

Co-authored-by: Isaac
---
 cmd/lakebox/api.go | 30 +++++++++++++++++++++++++++---
 cmd/lakebox/ssh.go | 25 +++++++++++++++++++++----
 2 files changed, 48 insertions(+), 7 deletions(-)

diff --git a/cmd/lakebox/api.go b/cmd/lakebox/api.go
index 4fb3af2630..acaeff47e8 100644
--- a/cmd/lakebox/api.go
+++ b/cmd/lakebox/api.go
@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"net/url"
 	"strings"
 	"time"
 
@@ -277,10 +278,24 @@ func (a *lakeboxAPI) delete(ctx context.Context, id string) error {
 
 // doRequest makes an authenticated HTTP request to the workspace.
 func (a *lakeboxAPI) doRequest(ctx context.Context, method, path string, body io.Reader) (*http.Response, error) {
-	host := strings.TrimRight(a.w.Config.Host, "/")
-	url := host + path
+	// The configured host may be just a hostname or may carry a workspace
+	// selector in the query (e.g. `https://dogfood.staging.databricks.com/?o=...`).
+	// Parse it so we can append the API path while preserving the query, and so
+	// we can pull the workspace ID out of `?o=<id>` when the SDK config doesn't
+	// carry it on a separate `workspace_id` field.
+	parsed, err := url.Parse(a.w.Config.Host)
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse host %q: %w", a.w.Config.Host, err)
+	}
+	wsid := a.w.Config.WorkspaceID
+	if wsid == "" {
+		if v := parsed.Query().Get("o"); v != "" {
+			wsid = v
+		}
+	}
+	parsed.Path = strings.TrimRight(parsed.Path, "/") + path
 
-	req, err := http.NewRequestWithContext(ctx, method, url, body)
+	req, err := http.NewRequestWithContext(ctx, method, parsed.String(), body)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create request: %w", err)
 	}
@@ -289,6 +304,15 @@ func (a *lakeboxAPI) doRequest(ctx context.Context, method, path string, body io
 		return nil, fmt.Errorf("failed to authenticate: %w", err)
 	}
 
+	// Multi-workspace gateways (e.g. dogfood.staging.databricks.com) need a
+	// workspace selector to route the request — without it the gateway can't
+	// scope the credential and rejects with "Credential was not sent or was of
+	// an unsupported type for this API". `?o=<id>` in the URL works as a
+	// fallback, but the explicit header is the well-defined contract.
+	if wsid != "" {
+		req.Header.Set("X-Databricks-Org-Id", wsid)
+	}
+
 	if body != nil {
 		req.Header.Set("Content-Type", "application/json")
 	}
diff --git a/cmd/lakebox/ssh.go b/cmd/lakebox/ssh.go
index 11297f2786..2a7db87a1b 100644
--- a/cmd/lakebox/ssh.go
+++ b/cmd/lakebox/ssh.go
@@ -5,16 +5,28 @@ import (
 	"os"
 	"os/exec"
 	"runtime"
+	"strings"
 
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/spf13/cobra"
 )
 
 const (
-	defaultGatewayHost = "uw2.dbrx.dev"
-	defaultGatewayPort = "2222"
+	defaultGatewayHost        = "uw2.dbrx.dev"
+	stagingDefaultGatewayHost = "uw2.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
+// uw2.s.dbrx.dev; everything else uses prod uw2.dbrx.dev.
+func resolveGatewayHost(workspaceHost string) string {
+	if strings.Contains(workspaceHost, ".staging.") {
+		return stagingDefaultGatewayHost
+	}
+	return defaultGatewayHost
+}
+
 func newSSHCommand() *cobra.Command {
 	var gatewayHost string
 	var gatewayPort string
@@ -98,13 +110,18 @@ Examples:
 				}
 			}
 
+			host := gatewayHost
+			if host == "" {
+				host = resolveGatewayHost(w.Config.Host)
+			}
+
 			s := spin(stderr, fmt.Sprintf("Connecting to %s…", bold(lakeboxID)))
 			s.ok(fmt.Sprintf("Connected to %s", bold(lakeboxID)))
-			return execSSHDirect(lakeboxID, gatewayHost, gatewayPort, keyPath, extraArgs)
+			return execSSHDirect(lakeboxID, host, gatewayPort, keyPath, extraArgs)
 		},
 	}
 
-	cmd.Flags().StringVar(&gatewayHost, "gateway", defaultGatewayHost, "Lakebox gateway hostname")
+	cmd.Flags().StringVar(&gatewayHost, "gateway", "", "Lakebox gateway hostname (auto-detected from profile if empty)")
 	cmd.Flags().StringVar(&gatewayPort, "port", defaultGatewayPort, "Lakebox gateway SSH port")
 
 	return cmd

From 86b42952029f8d7adc7462b9f943d9809cca9efc Mon Sep 17 00:00:00 2001
From: Pieter Noordhuis <pieter.noordhuis@databricks.com>
Date: Fri, 8 May 2026 17:02:35 +0200
Subject: [PATCH 15/36] lakebox: integrate as a 'databricks lakebox' subcommand
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Wire the cmd/lakebox tree from #4930 into the main CLI:

 - cmd/cmd.go registers lakebox.New() under the 'development' command
   group alongside bundle and sync.
 - cmd/fuzz_panic_test.go adds 'lakebox' to manualRoots so TestCountFuzz
   doesn't fuzz hand-written commands as if they were auto-generated.
 - cmd/lakebox tree: the original PR's standalone-CLI scaffolding is
   adapted for subcommand use — drop the auth-login hijacking and its
   helper exports, drop the 'last_profile' state field that only mattered
   when lakebox owned the whole CLI, switch PreRunE to root.MustWorkspaceClient
   directly, and update help text from 'lakebox foo' to
   'databricks lakebox foo' throughout.

Also conforms cmd/lakebox to project lint rules: env.UserHomeDir(ctx)
in place of os.UserHomeDir, errors.Is(err, fs.ErrNotExist) instead of
os.IsNotExist, atomic.Bool over sync.Once in the spinner gate, errors.New
for static error strings.

Co-authored-by: Isaac
---
 cmd/cmd.go              |  2 ++
 cmd/fuzz_panic_test.go  |  1 +
 cmd/lakebox/config.go   | 20 +++++++------
 cmd/lakebox/create.go   |  9 +++---
 cmd/lakebox/default.go  |  8 ++++--
 cmd/lakebox/delete.go   | 17 +++++------
 cmd/lakebox/lakebox.go  | 38 ++++++++-----------------
 cmd/lakebox/list.go     | 24 ++++++----------
 cmd/lakebox/register.go | 51 +++++++++++----------------------
 cmd/lakebox/ssh.go      | 39 ++++++++++++++------------
 cmd/lakebox/state.go    | 62 +++++++++++++++--------------------------
 cmd/lakebox/status.go   |  3 +-
 cmd/lakebox/ui.go       | 46 +++++++++++++++---------------
 13 files changed, 139 insertions(+), 181 deletions(-)

diff --git a/cmd/cmd.go b/cmd/cmd.go
index 014471f763..d8a8c09f04 100644
--- a/cmd/cmd.go
+++ b/cmd/cmd.go
@@ -17,6 +17,7 @@ import (
 	"github.com/databricks/cli/cmd/experimental"
 	"github.com/databricks/cli/cmd/fs"
 	"github.com/databricks/cli/cmd/labs"
+	"github.com/databricks/cli/cmd/lakebox"
 	"github.com/databricks/cli/cmd/pipelines"
 	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/cmd/selftest"
@@ -103,6 +104,7 @@ func New(ctx context.Context) *cobra.Command {
 	cli.AddCommand(configure.New())
 	cli.AddCommand(fs.New())
 	cli.AddCommand(labs.New(ctx))
+	cli.AddCommand(lakebox.New())
 	cli.AddCommand(sync.New())
 	cli.AddCommand(version.New())
 	cli.AddCommand(selftest.New())
diff --git a/cmd/fuzz_panic_test.go b/cmd/fuzz_panic_test.go
index 4fb5d5b9d3..e4037b4ef8 100644
--- a/cmd/fuzz_panic_test.go
+++ b/cmd/fuzz_panic_test.go
@@ -208,6 +208,7 @@ func isAutoGenerated(leaf leafCommand) bool {
 		"configure":    true,
 		"experimental": true,
 		"labs":         true,
+		"lakebox":      true,
 		"pipelines":    true,
 		"psql":         true,
 		"selftest":     true,
diff --git a/cmd/lakebox/config.go b/cmd/lakebox/config.go
index fe3b80ddf2..05fc157b04 100644
--- a/cmd/lakebox/config.go
+++ b/cmd/lakebox/config.go
@@ -1,9 +1,11 @@
 package lakebox
 
 import (
+	"errors"
 	"fmt"
 	"time"
 
+	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/spf13/cobra"
 )
@@ -39,17 +41,17 @@ Two knobs are independent — pass either or both:
                               this is on. Setting --idle-timeout to a
                               non-zero value in a later call clears
                               --no-autostop automatically. Sandbox still
-                              stops on explicit 'lakebox delete'.
+                              stops on explicit 'databricks lakebox delete'.
 
 Examples:
-  lakebox config happy-panda-1234 --idle-timeout 15m
-  lakebox config happy-panda-1234 --idle-timeout 1h30m
-  lakebox config happy-panda-1234 --idle-timeout 0           # clear, use default
-  lakebox config happy-panda-1234 --no-autostop                  # never auto-stop
-  lakebox config happy-panda-1234 --no-autostop=false            # back to timeout path
-  lakebox config happy-panda-1234 --idle-timeout 30m --no-autostop=false`,
+  databricks lakebox config happy-panda-1234 --idle-timeout 15m
+  databricks lakebox config happy-panda-1234 --idle-timeout 1h30m
+  databricks lakebox config happy-panda-1234 --idle-timeout 0           # clear, use default
+  databricks lakebox config happy-panda-1234 --no-autostop                  # never auto-stop
+  databricks lakebox config happy-panda-1234 --no-autostop=false            # back to timeout path
+  databricks lakebox config happy-panda-1234 --idle-timeout 30m --no-autostop=false`,
 		Args:    cobra.ExactArgs(1),
-		PreRunE: mustWorkspaceClient,
+		PreRunE: root.MustWorkspaceClient,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := cmd.Context()
 			w := cmdctx.WorkspaceClient(ctx)
@@ -76,7 +78,7 @@ Examples:
 			}
 
 			if idleSecs == nil && noAutostop == nil {
-				return fmt.Errorf("nothing to update — pass --idle-timeout and/or --no-autostop")
+				return errors.New("nothing to update — pass --idle-timeout and/or --no-autostop")
 			}
 
 			updated, err := api.update(ctx, id, idleSecs, noAutostop)
diff --git a/cmd/lakebox/create.go b/cmd/lakebox/create.go
index 096df26ce6..ee33ae7e08 100644
--- a/cmd/lakebox/create.go
+++ b/cmd/lakebox/create.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"os"
 
+	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/spf13/cobra"
 )
@@ -20,8 +21,8 @@ Creates a new personal development environment backed by a microVM.
 Blocks until the lakebox is running and prints the lakebox ID.
 
 Example:
-  lakebox create`,
-		PreRunE: mustWorkspaceClient,
+  databricks lakebox create`,
+		PreRunE: root.MustWorkspaceClient,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := cmd.Context()
 			w := cmdctx.WorkspaceClient(ctx)
@@ -52,7 +53,7 @@ Example:
 				profile = w.Config.Host
 			}
 
-			currentDefault := getDefault(profile)
+			currentDefault := getDefault(ctx, profile)
 			shouldSetDefault := currentDefault == ""
 			if !shouldSetDefault && currentDefault != "" {
 				if _, err := api.get(ctx, currentDefault); err != nil {
@@ -60,7 +61,7 @@ Example:
 				}
 			}
 			if shouldSetDefault {
-				if err := setDefault(profile, result.SandboxID); err != nil {
+				if err := setDefault(ctx, profile, result.SandboxID); err != nil {
 					warn(stderr, fmt.Sprintf("Could not save default: %v", err))
 				} else {
 					field(stderr, "default", result.SandboxID)
diff --git a/cmd/lakebox/default.go b/cmd/lakebox/default.go
index b632c5984a..cd96df172d 100644
--- a/cmd/lakebox/default.go
+++ b/cmd/lakebox/default.go
@@ -3,6 +3,7 @@ package lakebox
 import (
 	"fmt"
 
+	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/spf13/cobra"
 )
@@ -18,16 +19,17 @@ The default is stored locally in ~/.databricks/lakebox.json per profile.
 Example:
   databricks lakebox set-default happy-panda-1234`,
 		Args:    cobra.ExactArgs(1),
-		PreRunE: mustWorkspaceClient,
+		PreRunE: root.MustWorkspaceClient,
 		RunE: func(cmd *cobra.Command, args []string) error {
-			w := cmdctx.WorkspaceClient(cmd.Context())
+			ctx := cmd.Context()
+			w := cmdctx.WorkspaceClient(ctx)
 			profile := w.Config.Profile
 			if profile == "" {
 				profile = w.Config.Host
 			}
 
 			lakeboxID := args[0]
-			if err := setDefault(profile, lakeboxID); err != nil {
+			if err := setDefault(ctx, profile, lakeboxID); err != nil {
 				return fmt.Errorf("failed to set default: %w", err)
 			}
 			fmt.Fprintf(cmd.OutOrStdout(), "Default lakebox set to: %s\n", lakeboxID)
diff --git a/cmd/lakebox/delete.go b/cmd/lakebox/delete.go
index ba56e2a508..c95c8bb965 100644
--- a/cmd/lakebox/delete.go
+++ b/cmd/lakebox/delete.go
@@ -3,6 +3,7 @@ package lakebox
 import (
 	"fmt"
 
+	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/spf13/cobra"
 )
@@ -16,9 +17,9 @@ func newDeleteCommand() *cobra.Command {
 Permanently terminates and removes the specified lakebox.
 
 Example:
-  lakebox delete happy-panda-1234`,
+  databricks lakebox delete happy-panda-1234`,
 		Args:    cobra.ExactArgs(1),
-		PreRunE: mustWorkspaceClient,
+		PreRunE: root.MustWorkspaceClient,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := cmd.Context()
 			w := cmdctx.WorkspaceClient(ctx)
@@ -26,10 +27,10 @@ Example:
 			stderr := cmd.ErrOrStderr()
 
 			lakeboxID := args[0]
-			s := spin(stderr, fmt.Sprintf("Removing %s…", lakeboxID))
+			s := spin(stderr, "Removing "+lakeboxID+"…")
 
 			if err := api.delete(ctx, lakeboxID); err != nil {
-				s.fail(fmt.Sprintf("Failed to delete %s", lakeboxID))
+				s.fail("Failed to delete " + lakeboxID)
 				return fmt.Errorf("failed to delete lakebox %s: %w", lakeboxID, err)
 			}
 
@@ -37,11 +38,11 @@ Example:
 			if profile == "" {
 				profile = w.Config.Host
 			}
-			if getDefault(profile) == lakeboxID {
-				_ = clearDefault(profile)
-				s.ok(fmt.Sprintf("Removed %s %s", bold(lakeboxID), dim("(default cleared)")))
+			if getDefault(ctx, profile) == lakeboxID {
+				_ = clearDefault(ctx, profile)
+				s.ok("Removed " + bold(lakeboxID) + " " + dim("(default cleared)"))
 			} else {
-				s.ok(fmt.Sprintf("Removed %s", bold(lakeboxID)))
+				s.ok("Removed " + bold(lakeboxID))
 			}
 			return nil
 		},
diff --git a/cmd/lakebox/lakebox.go b/cmd/lakebox/lakebox.go
index 25a9b479e5..b6c2970760 100644
--- a/cmd/lakebox/lakebox.go
+++ b/cmd/lakebox/lakebox.go
@@ -1,14 +1,14 @@
 package lakebox
 
 import (
-	"github.com/databricks/cli/cmd/root"
 	"github.com/spf13/cobra"
 )
 
 func New() *cobra.Command {
 	cmd := &cobra.Command{
-		Use:   "lakebox",
-		Short: "Manage Databricks Lakebox environments",
+		Use:     "lakebox",
+		Short:   "Manage Databricks Lakebox environments",
+		GroupID: "development",
 		Long: `Manage Databricks Lakebox environments.
 
 Lakebox provides SSH-accessible development environments backed by
@@ -16,19 +16,17 @@ microVM isolation. Each lakebox is a personal sandbox with pre-installed
 tooling (Python, Node.js, Rust, Databricks CLI) and persistent storage.
 
 Getting started:
-  lakebox auth login --host https://...   # authenticate to Databricks workspace and lakebox service
-  lakebox ssh                             # SSH to your default lakebox
+  databricks auth login --host https://...   # authenticate to a Databricks workspace
+  databricks lakebox register                # generate and register an SSH key
+  databricks lakebox ssh                     # SSH to your default lakebox
 
 Common workflows:
-  lakebox ssh                             # SSH to your default lakebox
-  lakebox ssh my-project                  # SSH to a named lakebox
-  lakebox list                            # list your lakeboxes
-  lakebox create                          # create a new lakebox
-  lakebox delete my-project               # delete a lakebox
-  lakebox status my-project               # show lakebox status
-
-The CLI manages your ~/.ssh/config so you can also connect directly:
-  ssh my-project                          # after 'lakebox ssh'
+  databricks lakebox ssh                     # SSH to your default lakebox
+  databricks lakebox ssh my-project          # SSH to a named lakebox
+  databricks lakebox list                    # list your lakeboxes
+  databricks lakebox create                  # create a new lakebox
+  databricks lakebox delete my-project       # delete a lakebox
+  databricks lakebox status my-project       # show lakebox status
 `,
 	}
 
@@ -43,15 +41,3 @@ The CLI manages your ~/.ssh/config so you can also connect directly:
 
 	return cmd
 }
-
-// mustWorkspaceClient applies the saved last-login profile when the user
-// hasn't explicitly set --profile, then delegates to root.MustWorkspaceClient.
-func mustWorkspaceClient(cmd *cobra.Command, args []string) error {
-	profileFlag := cmd.Flag("profile")
-	if profileFlag != nil && !profileFlag.Changed {
-		if last := GetLastProfile(); last != "" {
-			_ = profileFlag.Value.Set(last)
-		}
-	}
-	return root.MustWorkspaceClient(cmd, args)
-}
diff --git a/cmd/lakebox/list.go b/cmd/lakebox/list.go
index f058524e7e..7c2fd62c88 100644
--- a/cmd/lakebox/list.go
+++ b/cmd/lakebox/list.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"strings"
 
+	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/spf13/cobra"
 )
@@ -21,9 +22,9 @@ Shows all lakeboxes associated with your account, including their
 current status and ID.
 
 Example:
-  lakebox list
-  lakebox list --json`,
-		PreRunE: mustWorkspaceClient,
+  databricks lakebox list
+  databricks lakebox list --json`,
+		PreRunE: root.MustWorkspaceClient,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := cmd.Context()
 			w := cmdctx.WorkspaceClient(ctx)
@@ -49,7 +50,7 @@ Example:
 			if profile == "" {
 				profile = w.Config.Host
 			}
-			defaultID := getDefault(profile)
+			defaultID := getDefault(ctx, profile)
 
 			out := cmd.OutOrStdout()
 
@@ -80,21 +81,12 @@ Example:
 					def = accent("*")
 				}
 				// Pad ID manually to avoid ANSI codes breaking alignment.
-				idPad := col - len(id)
-				if idPad < 0 {
-					idPad = 0
-				}
+				idPad := max(col-len(id), 0)
 				st := status(e.Status)
 				// Pad status to 10 visible chars.
-				stPad := 10 - len(e.Status)
-				if stPad < 0 {
-					stPad = 0
-				}
+				stPad := max(10-len(e.Status), 0)
 				as := e.autoStopLabel()
-				asPad := autostopCol - len(as)
-				if asPad < 0 {
-					asPad = 0
-				}
+				asPad := max(autostopCol-len(as), 0)
 				idStr := bold(id)
 				if strings.EqualFold(e.Status, "running") {
 					idStr = cyan + bo + id + rs
diff --git a/cmd/lakebox/register.go b/cmd/lakebox/register.go
index f3550d8e5d..0f4e0bc5b9 100644
--- a/cmd/lakebox/register.go
+++ b/cmd/lakebox/register.go
@@ -2,13 +2,15 @@ package lakebox
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"os"
 	"os/exec"
 	"path/filepath"
 
+	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdctx"
-	"github.com/databricks/databricks-sdk-go"
+	"github.com/databricks/cli/libs/env"
 	"github.com/spf13/cobra"
 )
 
@@ -24,25 +26,25 @@ This command:
 1. Generates an RSA SSH key at ~/.ssh/lakebox_rsa (if it doesn't exist)
 2. Registers the public key with the lakebox service
 
-After registration, 'lakebox ssh' will use this key automatically.
+After registration, 'databricks lakebox ssh' will use this key automatically.
 Run this once per machine.
 
 Example:
-  lakebox register`,
-		PreRunE: mustWorkspaceClient,
+  databricks lakebox register`,
+		PreRunE: root.MustWorkspaceClient,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := cmd.Context()
 			w := cmdctx.WorkspaceClient(ctx)
 			api := newLakeboxAPI(w)
 
-			keyPath, generated, err := ensureLakeboxKey()
+			keyPath, generated, err := ensureLakeboxKey(ctx)
 			if err != nil {
 				return fmt.Errorf("failed to ensure lakebox SSH key: %w", err)
 			}
 
 			stderr := cmd.ErrOrStderr()
 			if generated {
-				ok(stderr, fmt.Sprintf("Generated SSH key at %s", dim(keyPath)))
+				ok(stderr, "Generated SSH key at "+dim(keyPath))
 			} else {
 				field(stderr, "key", keyPath)
 			}
@@ -60,7 +62,7 @@ Example:
 			s.ok("SSH key registered")
 
 			blank(stderr)
-			fmt.Fprintf(stderr, "  Run %s to connect.\n\n", bold("lakebox ssh"))
+			fmt.Fprintf(stderr, "  Run %s to connect.\n\n", bold("databricks lakebox ssh"))
 			return nil
 		},
 	}
@@ -69,8 +71,8 @@ Example:
 }
 
 // lakeboxKeyPath returns the path to the dedicated lakebox SSH key.
-func lakeboxKeyPath() (string, error) {
-	homeDir, err := os.UserHomeDir()
+func lakeboxKeyPath(ctx context.Context) (string, error) {
+	homeDir, err := env.UserHomeDir(ctx)
 	if err != nil {
 		return "", err
 	}
@@ -79,8 +81,8 @@ func lakeboxKeyPath() (string, error) {
 
 // ensureLakeboxKey returns the path to the lakebox SSH key, generating it if
 // it doesn't exist. Returns (path, wasGenerated, error).
-func ensureLakeboxKey() (string, bool, error) {
-	keyPath, err := lakeboxKeyPath()
+func ensureLakeboxKey(ctx context.Context) (string, bool, error) {
+	keyPath, err := lakeboxKeyPath(ctx)
 	if err != nil {
 		return "", false, err
 	}
@@ -91,16 +93,16 @@ func ensureLakeboxKey() (string, bool, error) {
 
 	// Check that ssh-keygen is available before trying to generate.
 	if _, err := exec.LookPath("ssh-keygen"); err != nil {
-		return "", false, fmt.Errorf(
+		return "", false, errors.New(
 			"ssh-keygen not found in PATH.\n" +
-				"Please install OpenSSH and run 'lakebox register' again.\n" +
+				"Please install OpenSSH and run 'databricks lakebox register' again.\n" +
 				"  macOS:   brew install openssh\n" +
 				"  Ubuntu:  sudo apt install openssh-client\n" +
 				"  Windows: install Git for Windows (includes ssh-keygen)")
 	}
 
 	sshDir := filepath.Dir(keyPath)
-	if err := os.MkdirAll(sshDir, 0700); err != nil {
+	if err := os.MkdirAll(sshDir, 0o700); err != nil {
 		return "", false, fmt.Errorf("failed to create %s: %w", sshDir, err)
 	}
 
@@ -114,24 +116,3 @@ func ensureLakeboxKey() (string, bool, error) {
 
 	return keyPath, true, nil
 }
-
-// EnsureAndReadKey generates the lakebox SSH key if needed and returns
-// (keyPath, publicKeyContent, error). Exported for use by the auth login hook.
-func EnsureAndReadKey() (string, string, error) {
-	keyPath, _, err := ensureLakeboxKey()
-	if err != nil {
-		return "", "", err
-	}
-	pubKeyData, err := os.ReadFile(keyPath + ".pub")
-	if err != nil {
-		return "", "", fmt.Errorf("failed to read public key %s.pub: %w", keyPath, err)
-	}
-	return keyPath, string(pubKeyData), nil
-}
-
-// RegisterKey registers a public key with the lakebox API. Exported for use
-// by the auth login hook.
-func RegisterKey(ctx context.Context, w *databricks.WorkspaceClient, pubKey string) error {
-	api := newLakeboxAPI(w)
-	return api.registerKey(ctx, pubKey)
-}
diff --git a/cmd/lakebox/ssh.go b/cmd/lakebox/ssh.go
index 2a7db87a1b..6ac716d158 100644
--- a/cmd/lakebox/ssh.go
+++ b/cmd/lakebox/ssh.go
@@ -1,12 +1,15 @@
 package lakebox
 
 import (
+	"errors"
 	"fmt"
+	"io/fs"
 	"os"
 	"os/exec"
 	"runtime"
 	"strings"
 
+	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/spf13/cobra"
 )
@@ -41,13 +44,13 @@ after -- are passed directly to the ssh process. This lets you run
 remote commands, set up port forwarding, or pass any other ssh flags.
 
 Examples:
-  lakebox ssh                                  # interactive shell on default lakebox
-  lakebox ssh happy-panda-1234                 # interactive shell on specific lakebox
-  lakebox ssh -- ls -la /home                  # run command on default lakebox
-  lakebox ssh happy-panda-1234 -- cat /etc/os-release  # run command on specific lakebox
-  lakebox ssh -- -L 8080:localhost:8080        # port forwarding on default lakebox`,
+  databricks lakebox ssh                                  # interactive shell on default lakebox
+  databricks lakebox ssh happy-panda-1234                 # interactive shell on specific lakebox
+  databricks lakebox ssh -- ls -la /home                  # run command on default lakebox
+  databricks lakebox ssh happy-panda-1234 -- cat /etc/os-release  # run command on specific lakebox
+  databricks lakebox ssh -- -L 8080:localhost:8080        # port forwarding on default lakebox`,
 		Args:    cobra.ArbitraryArgs,
-		PreRunE: mustWorkspaceClient,
+		PreRunE: root.MustWorkspaceClient,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := cmd.Context()
 			w := cmdctx.WorkspaceClient(ctx)
@@ -58,12 +61,12 @@ Examples:
 			}
 
 			// Use the dedicated lakebox SSH key.
-			keyPath, err := lakeboxKeyPath()
+			keyPath, err := lakeboxKeyPath(ctx)
 			if err != nil {
 				return fmt.Errorf("failed to determine lakebox key path: %w", err)
 			}
-			if _, err := os.Stat(keyPath); os.IsNotExist(err) {
-				return fmt.Errorf("lakebox SSH key not found at %s — run 'lakebox register' first", keyPath)
+			if _, err := os.Stat(keyPath); errors.Is(err, fs.ErrNotExist) {
+				return fmt.Errorf("lakebox SSH key not found at %s — run 'databricks lakebox register' first", keyPath)
 			}
 			stderr := cmd.ErrOrStderr()
 
@@ -72,21 +75,21 @@ Examples:
 			var lakeboxID string
 			var extraArgs []string
 
-			dashAt := cmd.ArgsLenAtDash()
-			if dashAt == -1 {
+			switch dashAt := cmd.ArgsLenAtDash(); dashAt {
+			case -1:
 				if len(args) > 0 {
 					lakeboxID = args[0]
 				}
-			} else if dashAt == 0 {
+			case 0:
 				extraArgs = args[dashAt:]
-			} else {
+			default:
 				lakeboxID = args[0]
 				extraArgs = args[dashAt:]
 			}
 
 			// Determine lakebox ID if not explicit.
 			if lakeboxID == "" {
-				if def := getDefault(profile); def != "" {
+				if def := getDefault(ctx, profile); def != "" {
 					lakeboxID = def
 				} else {
 					api := newLakeboxAPI(w)
@@ -102,9 +105,9 @@ Examples:
 						return fmt.Errorf("failed to create lakebox: %w", err)
 					}
 					lakeboxID = result.SandboxID
-					s.ok(fmt.Sprintf("Lakebox %s ready", bold(lakeboxID)))
+					s.ok("Lakebox " + bold(lakeboxID) + " ready")
 
-					if err := setDefault(profile, lakeboxID); err != nil {
+					if err := setDefault(ctx, profile, lakeboxID); err != nil {
 						warn(stderr, fmt.Sprintf("Could not save default: %v", err))
 					}
 				}
@@ -115,8 +118,8 @@ Examples:
 				host = resolveGatewayHost(w.Config.Host)
 			}
 
-			s := spin(stderr, fmt.Sprintf("Connecting to %s…", bold(lakeboxID)))
-			s.ok(fmt.Sprintf("Connected to %s", bold(lakeboxID)))
+			s := spin(stderr, "Connecting to "+bold(lakeboxID)+"…")
+			s.ok("Connected to " + bold(lakeboxID))
 			return execSSHDirect(lakeboxID, host, gatewayPort, keyPath, extraArgs)
 		},
 	}
diff --git a/cmd/lakebox/state.go b/cmd/lakebox/state.go
index b84b5b16e1..87cc96e78f 100644
--- a/cmd/lakebox/state.go
+++ b/cmd/lakebox/state.go
@@ -1,10 +1,15 @@
 package lakebox
 
 import (
+	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
+	"io/fs"
 	"os"
 	"path/filepath"
+
+	"github.com/databricks/cli/libs/env"
 )
 
 // stateFile stores per-profile lakebox defaults on the local filesystem.
@@ -12,26 +17,24 @@ import (
 type stateFile struct {
 	// Profile name → default lakebox ID.
 	Defaults map[string]string `json:"defaults"`
-	// Last profile used with 'lakebox auth login'.
-	LastProfile string `json:"last_profile,omitempty"`
 }
 
-func stateFilePath() (string, error) {
-	home, err := os.UserHomeDir()
+func stateFilePath(ctx context.Context) (string, error) {
+	home, err := env.UserHomeDir(ctx)
 	if err != nil {
 		return "", err
 	}
 	return filepath.Join(home, ".databricks", "lakebox.json"), nil
 }
 
-func loadState() (*stateFile, error) {
-	path, err := stateFilePath()
+func loadState(ctx context.Context) (*stateFile, error) {
+	path, err := stateFilePath(ctx)
 	if err != nil {
 		return nil, err
 	}
 
 	data, err := os.ReadFile(path)
-	if os.IsNotExist(err) {
+	if errors.Is(err, fs.ErrNotExist) {
 		return &stateFile{Defaults: make(map[string]string)}, nil
 	}
 	if err != nil {
@@ -40,7 +43,7 @@ func loadState() (*stateFile, error) {
 
 	var state stateFile
 	if err := json.Unmarshal(data, &state); err != nil {
-		return &stateFile{Defaults: make(map[string]string)}, nil
+		return nil, fmt.Errorf("failed to parse %s: %w", path, err)
 	}
 	if state.Defaults == nil {
 		state.Defaults = make(map[string]string)
@@ -48,13 +51,13 @@ func loadState() (*stateFile, error) {
 	return &state, nil
 }
 
-func saveState(state *stateFile) error {
-	path, err := stateFilePath()
+func saveState(ctx context.Context, state *stateFile) error {
+	path, err := stateFilePath(ctx)
 	if err != nil {
 		return err
 	}
 
-	if err := os.MkdirAll(filepath.Dir(path), 0700); err != nil {
+	if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
 		return err
 	}
 
@@ -62,50 +65,31 @@ func saveState(state *stateFile) error {
 	if err != nil {
 		return err
 	}
-	return os.WriteFile(path, data, 0600)
+	return os.WriteFile(path, data, 0o600)
 }
 
-func getDefault(profile string) string {
-	state, err := loadState()
+func getDefault(ctx context.Context, profile string) string {
+	state, err := loadState(ctx)
 	if err != nil {
 		return ""
 	}
 	return state.Defaults[profile]
 }
 
-func setDefault(profile, lakeboxID string) error {
-	state, err := loadState()
+func setDefault(ctx context.Context, profile, lakeboxID string) error {
+	state, err := loadState(ctx)
 	if err != nil {
 		return err
 	}
 	state.Defaults[profile] = lakeboxID
-	return saveState(state)
-}
-
-// GetLastProfile returns the profile saved by the most recent 'lakebox auth login'.
-func GetLastProfile() string {
-	state, err := loadState()
-	if err != nil {
-		return ""
-	}
-	return state.LastProfile
-}
-
-// SetLastProfile persists the profile used during 'lakebox auth login'.
-func SetLastProfile(profile string) error {
-	state, err := loadState()
-	if err != nil {
-		return err
-	}
-	state.LastProfile = profile
-	return saveState(state)
+	return saveState(ctx, state)
 }
 
-func clearDefault(profile string) error {
-	state, err := loadState()
+func clearDefault(ctx context.Context, profile string) error {
+	state, err := loadState(ctx)
 	if err != nil {
 		return err
 	}
 	delete(state.Defaults, profile)
-	return saveState(state)
+	return saveState(ctx, state)
 }
diff --git a/cmd/lakebox/status.go b/cmd/lakebox/status.go
index f5df1ee4a4..7050c53ff6 100644
--- a/cmd/lakebox/status.go
+++ b/cmd/lakebox/status.go
@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"fmt"
 
+	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/spf13/cobra"
 )
@@ -20,7 +21,7 @@ Example:
   lakebox status happy-panda-1234
   lakebox status happy-panda-1234 --json`,
 		Args:    cobra.ExactArgs(1),
-		PreRunE: mustWorkspaceClient,
+		PreRunE: root.MustWorkspaceClient,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := cmd.Context()
 			w := cmdctx.WorkspaceClient(ctx)
diff --git a/cmd/lakebox/ui.go b/cmd/lakebox/ui.go
index 2eab33310c..de1aa2cf91 100644
--- a/cmd/lakebox/ui.go
+++ b/cmd/lakebox/ui.go
@@ -5,7 +5,7 @@ import (
 	"io"
 	"os"
 	"strings"
-	"sync"
+	"sync/atomic"
 	"time"
 )
 
@@ -30,11 +30,11 @@ func isTTY(w io.Writer) bool {
 
 // spinner shows a braille spinner like Claude Code.
 type spinner struct {
-	w       io.Writer
-	msg     string
-	done    chan struct{}
-	once    sync.Once
-	started time.Time
+	w        io.Writer
+	msg      string
+	done     chan struct{}
+	finished atomic.Bool
+	started  time.Time
 }
 
 func spin(w io.Writer, msg string) *spinner {
@@ -68,25 +68,27 @@ func (s *spinner) run() {
 }
 
 func (s *spinner) ok(msg string) {
-	s.once.Do(func() {
-		close(s.done)
-		if isTTY(s.w) {
-			fmt.Fprintf(s.w, "\r\033[K  %s✓%s %s\n", cyan, rs, msg)
-		} else {
-			fmt.Fprintf(s.w, "✓ %s\n", msg)
-		}
-	})
+	if !s.finished.CompareAndSwap(false, true) {
+		return
+	}
+	close(s.done)
+	if isTTY(s.w) {
+		fmt.Fprintf(s.w, "\r\033[K  %s✓%s %s\n", cyan, rs, msg)
+	} else {
+		fmt.Fprintf(s.w, "✓ %s\n", msg)
+	}
 }
 
 func (s *spinner) fail(msg string) {
-	s.once.Do(func() {
-		close(s.done)
-		if isTTY(s.w) {
-			fmt.Fprintf(s.w, "\r\033[K  %s✗%s %s\n", cyan, rs, msg)
-		} else {
-			fmt.Fprintf(s.w, "✗ %s\n", msg)
-		}
-	})
+	if !s.finished.CompareAndSwap(false, true) {
+		return
+	}
+	close(s.done)
+	if isTTY(s.w) {
+		fmt.Fprintf(s.w, "\r\033[K  %s✗%s %s\n", cyan, rs, msg)
+	} else {
+		fmt.Fprintf(s.w, "✗ %s\n", msg)
+	}
 }
 
 // --- Consistent output primitives ---

From ea75d2c9b6a3f9f78dcc5790be145d9c5601067d Mon Sep 17 00:00:00 2001
From: Pieter Noordhuis <pieter.noordhuis@databricks.com>
Date: Wed, 6 May 2026 14:36:55 +0200
Subject: [PATCH 16/36] lakebox: rewrite ui.go on top of cmdio
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Replace the hand-rolled braille spinner, TTY detection, and stderr
plumbing with the existing cmdio facilities:

- spin(ctx, msg) wraps cmdio.NewSpinner — capability-aware, runs through
  the same Bubble Tea program slot as other CLI spinners. ok/fail markers
  are logged via cmdio.LogString after Close.
- ok(ctx, ...) and warn(ctx, ...) are now ctx-based and route to stderr
  through cmdio rather than taking a writer. Call sites drop their
  cmd.ErrOrStderr() locals where they were only used for these helpers.
- field/blank still take an io.Writer because callers need to target
  stdout for structured output (list, status, config).

Drops the local isTTY, atomic.Bool spinner gate, and ticker goroutine.

Co-authored-by: Isaac
---
 cmd/lakebox/create.go   |  11 ++--
 cmd/lakebox/delete.go   |   3 +-
 cmd/lakebox/register.go |   4 +-
 cmd/lakebox/ssh.go      |   7 +--
 cmd/lakebox/ui.go       | 129 +++++++++++-----------------------------
 5 files changed, 47 insertions(+), 107 deletions(-)

diff --git a/cmd/lakebox/create.go b/cmd/lakebox/create.go
index ee33ae7e08..5e1419a51d 100644
--- a/cmd/lakebox/create.go
+++ b/cmd/lakebox/create.go
@@ -27,7 +27,6 @@ Example:
 			ctx := cmd.Context()
 			w := cmdctx.WorkspaceClient(ctx)
 			api := newLakeboxAPI(w)
-			stderr := cmd.ErrOrStderr()
 
 			var publicKey string
 			if publicKeyFile != "" {
@@ -38,7 +37,7 @@ Example:
 				publicKey = string(data)
 			}
 
-			s := spin(stderr, "Provisioning your lakebox…")
+			s := spin(ctx, "Provisioning your lakebox…")
 
 			result, err := api.create(ctx, publicKey)
 			if err != nil {
@@ -46,7 +45,7 @@ Example:
 				return fmt.Errorf("failed to create lakebox: %w", err)
 			}
 
-			s.ok(fmt.Sprintf("Lakebox %s is %s", bold(result.SandboxID), status(result.Status)))
+			s.ok("Lakebox " + bold(result.SandboxID) + " is " + status(result.Status))
 
 			profile := w.Config.Profile
 			if profile == "" {
@@ -62,13 +61,13 @@ Example:
 			}
 			if shouldSetDefault {
 				if err := setDefault(ctx, profile, result.SandboxID); err != nil {
-					warn(stderr, fmt.Sprintf("Could not save default: %v", err))
+					warn(ctx, fmt.Sprintf("Could not save default: %v", err))
 				} else {
-					field(stderr, "default", result.SandboxID)
+					field(cmd.ErrOrStderr(), "default", result.SandboxID)
 				}
 			}
 
-			blank(stderr)
+			blank(cmd.ErrOrStderr())
 			fmt.Fprintln(cmd.OutOrStdout(), result.SandboxID)
 			return nil
 		},
diff --git a/cmd/lakebox/delete.go b/cmd/lakebox/delete.go
index c95c8bb965..f65382b905 100644
--- a/cmd/lakebox/delete.go
+++ b/cmd/lakebox/delete.go
@@ -24,10 +24,9 @@ Example:
 			ctx := cmd.Context()
 			w := cmdctx.WorkspaceClient(ctx)
 			api := newLakeboxAPI(w)
-			stderr := cmd.ErrOrStderr()
 
 			lakeboxID := args[0]
-			s := spin(stderr, "Removing "+lakeboxID+"…")
+			s := spin(ctx, "Removing "+lakeboxID+"…")
 
 			if err := api.delete(ctx, lakeboxID); err != nil {
 				s.fail("Failed to delete " + lakeboxID)
diff --git a/cmd/lakebox/register.go b/cmd/lakebox/register.go
index 0f4e0bc5b9..8fdf42e2a7 100644
--- a/cmd/lakebox/register.go
+++ b/cmd/lakebox/register.go
@@ -44,7 +44,7 @@ Example:
 
 			stderr := cmd.ErrOrStderr()
 			if generated {
-				ok(stderr, "Generated SSH key at "+dim(keyPath))
+				ok(ctx, "Generated SSH key at "+dim(keyPath))
 			} else {
 				field(stderr, "key", keyPath)
 			}
@@ -54,7 +54,7 @@ Example:
 				return fmt.Errorf("failed to read public key %s.pub: %w", keyPath, err)
 			}
 
-			s := spin(stderr, "Registering key…")
+			s := spin(ctx, "Registering key…")
 			if err := api.registerKey(ctx, string(pubKeyData)); err != nil {
 				s.fail("Failed to register key")
 				return fmt.Errorf("failed to register key: %w", err)
diff --git a/cmd/lakebox/ssh.go b/cmd/lakebox/ssh.go
index 6ac716d158..bed21de997 100644
--- a/cmd/lakebox/ssh.go
+++ b/cmd/lakebox/ssh.go
@@ -68,7 +68,6 @@ Examples:
 			if _, err := os.Stat(keyPath); errors.Is(err, fs.ErrNotExist) {
 				return fmt.Errorf("lakebox SSH key not found at %s — run 'databricks lakebox register' first", keyPath)
 			}
-			stderr := cmd.ErrOrStderr()
 
 			// Parse args: everything before -- is the optional lakebox ID,
 			// everything after -- is passed through to ssh.
@@ -98,7 +97,7 @@ Examples:
 						return fmt.Errorf("failed to read public key %s.pub: %w", keyPath, err)
 					}
 
-					s := spin(stderr, "Provisioning your lakebox…")
+					s := spin(ctx, "Provisioning your lakebox…")
 					result, err := api.create(ctx, string(pubKeyData))
 					if err != nil {
 						s.fail("Failed to create lakebox")
@@ -108,7 +107,7 @@ Examples:
 					s.ok("Lakebox " + bold(lakeboxID) + " ready")
 
 					if err := setDefault(ctx, profile, lakeboxID); err != nil {
-						warn(stderr, fmt.Sprintf("Could not save default: %v", err))
+						warn(ctx, fmt.Sprintf("Could not save default: %v", err))
 					}
 				}
 			}
@@ -118,7 +117,7 @@ Examples:
 				host = resolveGatewayHost(w.Config.Host)
 			}
 
-			s := spin(stderr, "Connecting to "+bold(lakeboxID)+"…")
+			s := spin(ctx, "Connecting to "+bold(lakeboxID)+"…")
 			s.ok("Connected to " + bold(lakeboxID))
 			return execSSHDirect(lakeboxID, host, gatewayPort, keyPath, extraArgs)
 		},
diff --git a/cmd/lakebox/ui.go b/cmd/lakebox/ui.go
index de1aa2cf91..3d722986c3 100644
--- a/cmd/lakebox/ui.go
+++ b/cmd/lakebox/ui.go
@@ -1,15 +1,18 @@
 package lakebox
 
 import (
+	"context"
 	"fmt"
 	"io"
-	"os"
 	"strings"
-	"sync/atomic"
-	"time"
+
+	"github.com/databricks/cli/libs/cmdio"
 )
 
-// Single accent color throughout. Bold for emphasis. Dim for metadata.
+// ANSI escapes for inline highlighting. cmdio handles terminal capability
+// detection for the spinner, so we don't gate these on TTY here — strings
+// piped to a non-terminal still carry the codes, matching the behavior of
+// other CLI commands that call bold/dim helpers.
 const (
 	rs   = "\033[0m"  // reset
 	bo   = "\033[1m"  // bold
@@ -17,83 +20,34 @@ const (
 	cyan = "\033[36m" // accent
 )
 
-func isTTY(w io.Writer) bool {
-	if f, ok := w.(*os.File); ok {
-		fi, err := f.Stat()
-		if err != nil {
-			return false
-		}
-		return fi.Mode()&os.ModeCharDevice != 0
-	}
-	return false
-}
-
-// spinner shows a braille spinner like Claude Code.
+// spinner wraps cmdio.NewSpinner with terminal ok/fail markers. After the
+// first call to ok or fail, the spinner is closed and a final line is logged
+// to stderr; subsequent calls are no-ops.
 type spinner struct {
-	w        io.Writer
-	msg      string
-	done     chan struct{}
-	finished atomic.Bool
-	started  time.Time
+	ctx      context.Context
+	close    func()
+	finished bool
 }
 
-func spin(w io.Writer, msg string) *spinner {
-	s := &spinner{w: w, msg: msg, done: make(chan struct{}), started: time.Now()}
-	if isTTY(w) {
-		go s.run()
-	} else {
-		fmt.Fprintf(w, "* %s\n", msg)
-	}
-	return s
+func spin(ctx context.Context, msg string) *spinner {
+	sp := cmdio.NewSpinner(ctx)
+	sp.Update(msg)
+	return &spinner{ctx: ctx, close: sp.Close}
 }
 
-func (s *spinner) run() {
-	frames := []string{"⣾", "⣽", "⣻", "⢿", "⡿", "⣟", "⣯", "⣷"}
-	i := 0
-	ticker := time.NewTicker(80 * time.Millisecond)
-	defer ticker.Stop()
-	for {
-		select {
-		case <-s.done:
-			return
-		case <-ticker.C:
-			elapsed := time.Since(s.started).Truncate(time.Second)
-			fmt.Fprintf(s.w, "\r  %s%s%s %s%s%s %s(%s)%s  ",
-				cyan, frames[i%len(frames)], rs,
-				bo, s.msg, rs,
-				dm, elapsed, rs)
-			i++
-		}
-	}
-}
+func (s *spinner) ok(msg string)   { s.done("✓", msg) }
+func (s *spinner) fail(msg string) { s.done("✗", msg) }
 
-func (s *spinner) ok(msg string) {
-	if !s.finished.CompareAndSwap(false, true) {
+func (s *spinner) done(mark, msg string) {
+	if s.finished {
 		return
 	}
-	close(s.done)
-	if isTTY(s.w) {
-		fmt.Fprintf(s.w, "\r\033[K  %s✓%s %s\n", cyan, rs, msg)
-	} else {
-		fmt.Fprintf(s.w, "✓ %s\n", msg)
-	}
+	s.finished = true
+	s.close()
+	cmdio.LogString(s.ctx, "  "+cyan+mark+rs+" "+msg)
 }
 
-func (s *spinner) fail(msg string) {
-	if !s.finished.CompareAndSwap(false, true) {
-		return
-	}
-	close(s.done)
-	if isTTY(s.w) {
-		fmt.Fprintf(s.w, "\r\033[K  %s✗%s %s\n", cyan, rs, msg)
-	} else {
-		fmt.Fprintf(s.w, "✗ %s\n", msg)
-	}
-}
-
-// --- Consistent output primitives ---
-
-// status formats a status string with the accent color.
+// status formats a lakebox lifecycle status with the accent color.
 func status(s string) string {
 	switch strings.ToLower(s) {
 	case "running":
@@ -107,37 +61,26 @@ func status(s string) string {
 	}
 }
 
-// field prints "  label  value"
+// field prints "  label  value" to w.
 func field(w io.Writer, label, value string) {
 	fmt.Fprintf(w, "  %s%-10s%s %s\n", dm, label, rs, value)
 }
 
-// ok prints "  ✓ message"
-func ok(w io.Writer, msg string) {
-	fmt.Fprintf(w, "  %s✓%s %s\n", cyan, rs, msg)
+// ok prints "  ✓ message" to stderr via the cmdio context.
+func ok(ctx context.Context, msg string) {
+	cmdio.LogString(ctx, "  "+cyan+"✓"+rs+" "+msg)
 }
 
-// warn prints "  ! message"
-func warn(w io.Writer, msg string) {
-	fmt.Fprintf(w, "  %s!%s %s\n", cyan, rs, msg)
+// warn prints "  ! message" to stderr via the cmdio context.
+func warn(ctx context.Context, msg string) {
+	cmdio.LogString(ctx, "  "+cyan+"!"+rs+" "+msg)
 }
 
-// blank prints an empty line.
+// blank prints an empty line to w.
 func blank(w io.Writer) {
 	fmt.Fprintln(w)
 }
 
-// accent wraps text in the accent color.
-func accent(s string) string {
-	return cyan + s + rs
-}
-
-// bold wraps text in bold.
-func bold(s string) string {
-	return bo + s + rs
-}
-
-// dim wraps text in dim.
-func dim(s string) string {
-	return dm + s + rs
-}
+func accent(s string) string { return cyan + s + rs }
+func bold(s string) string   { return bo + s + rs }
+func dim(s string) string    { return dm + s + rs }

From 49cdfc3587ed39dddfdf4dd251913ddfd10fe5bd Mon Sep 17 00:00:00 2001
From: Pieter Noordhuis <pieter.noordhuis@databricks.com>
Date: Wed, 6 May 2026 15:47:04 +0200
Subject: [PATCH 17/36] lakebox: replace local ANSI consts with cmdio color
 helpers

Drops the cyan/bold/dim/reset constants and the local accent/bold/dim
wrappers in favor of cmdio.Cyan and cmdio.HiBlack, which respect the
SupportsStdoutColor capability check. Bold-for-emphasis is folded into
Cyan since cmdio does not expose a Go-level Bold helper today; visually
this means lakebox IDs and emphasized command names render in cyan
rather than uncolored bold, consistent with the rest of the CLI.

field/status now take a context so they can call cmdio.HiBlack /
cmdio.Cyan; their writer parameter stays for callers that target stdout.

Co-authored-by: Isaac
---
 cmd/lakebox/config.go   |  5 +++--
 cmd/lakebox/create.go   |  5 +++--
 cmd/lakebox/delete.go   |  5 +++--
 cmd/lakebox/list.go     | 25 ++++++++++++-------------
 cmd/lakebox/register.go |  7 ++++---
 cmd/lakebox/ssh.go      |  7 ++++---
 cmd/lakebox/status.go   |  9 +++++----
 cmd/lakebox/ui.go       | 39 ++++++++++++---------------------------
 8 files changed, 46 insertions(+), 56 deletions(-)

diff --git a/cmd/lakebox/config.go b/cmd/lakebox/config.go
index 05fc157b04..f3da157e00 100644
--- a/cmd/lakebox/config.go
+++ b/cmd/lakebox/config.go
@@ -7,6 +7,7 @@ import (
 
 	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdctx"
+	"github.com/databricks/cli/libs/cmdio"
 	"github.com/spf13/cobra"
 )
 
@@ -87,8 +88,8 @@ Examples:
 			}
 
 			blank(out)
-			field(out, "id", bold(updated.SandboxID))
-			field(out, "autostop", dim(updated.autoStopLabel()))
+			field(ctx, out, "id", cmdio.Cyan(ctx, updated.SandboxID))
+			field(ctx, out, "autostop", cmdio.HiBlack(ctx, updated.autoStopLabel()))
 			blank(out)
 			return nil
 		},
diff --git a/cmd/lakebox/create.go b/cmd/lakebox/create.go
index 5e1419a51d..62bcb3085d 100644
--- a/cmd/lakebox/create.go
+++ b/cmd/lakebox/create.go
@@ -6,6 +6,7 @@ import (
 
 	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdctx"
+	"github.com/databricks/cli/libs/cmdio"
 	"github.com/spf13/cobra"
 )
 
@@ -45,7 +46,7 @@ Example:
 				return fmt.Errorf("failed to create lakebox: %w", err)
 			}
 
-			s.ok("Lakebox " + bold(result.SandboxID) + " is " + status(result.Status))
+			s.ok("Lakebox " + cmdio.Cyan(ctx, result.SandboxID) + " is " + status(ctx, result.Status))
 
 			profile := w.Config.Profile
 			if profile == "" {
@@ -63,7 +64,7 @@ Example:
 				if err := setDefault(ctx, profile, result.SandboxID); err != nil {
 					warn(ctx, fmt.Sprintf("Could not save default: %v", err))
 				} else {
-					field(cmd.ErrOrStderr(), "default", result.SandboxID)
+					field(ctx, cmd.ErrOrStderr(), "default", result.SandboxID)
 				}
 			}
 
diff --git a/cmd/lakebox/delete.go b/cmd/lakebox/delete.go
index f65382b905..54d7a59e15 100644
--- a/cmd/lakebox/delete.go
+++ b/cmd/lakebox/delete.go
@@ -5,6 +5,7 @@ import (
 
 	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdctx"
+	"github.com/databricks/cli/libs/cmdio"
 	"github.com/spf13/cobra"
 )
 
@@ -39,9 +40,9 @@ Example:
 			}
 			if getDefault(ctx, profile) == lakeboxID {
 				_ = clearDefault(ctx, profile)
-				s.ok("Removed " + bold(lakeboxID) + " " + dim("(default cleared)"))
+				s.ok("Removed " + cmdio.Cyan(ctx, lakeboxID) + " " + cmdio.HiBlack(ctx, "(default cleared)"))
 			} else {
-				s.ok("Removed " + bold(lakeboxID))
+				s.ok("Removed " + cmdio.Cyan(ctx, lakeboxID))
 			}
 			return nil
 		},
diff --git a/cmd/lakebox/list.go b/cmd/lakebox/list.go
index 7c2fd62c88..d3d03f77fe 100644
--- a/cmd/lakebox/list.go
+++ b/cmd/lakebox/list.go
@@ -7,6 +7,7 @@ import (
 
 	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdctx"
+	"github.com/databricks/cli/libs/cmdio"
 	"github.com/spf13/cobra"
 )
 
@@ -42,7 +43,7 @@ Example:
 			}
 
 			if len(entries) == 0 {
-				fmt.Fprintf(cmd.ErrOrStderr(), "  %sNo lakeboxes found.%s\n", dm, rs)
+				fmt.Fprintf(cmd.ErrOrStderr(), "  %s\n", cmdio.HiBlack(ctx, "No lakeboxes found."))
 				return nil
 			}
 
@@ -70,31 +71,29 @@ Example:
 			autostopCol += 2
 
 			blank(out)
-			fmt.Fprintf(out, "  %s%-*s  %-10s  %-*s  %s%s\n",
-				dm, col, "ID", "STATUS", autostopCol, "AUTOSTOP", "DEFAULT", rs)
-			fmt.Fprintf(out, "  %s%s%s\n", dm, strings.Repeat("─", col+10+autostopCol+12), rs)
+			header := fmt.Sprintf("%-*s  %-10s  %-*s  %s",
+				col, "ID", "STATUS", autostopCol, "AUTOSTOP", "DEFAULT")
+			fmt.Fprintf(out, "  %s\n", cmdio.HiBlack(ctx, header))
+			fmt.Fprintf(out, "  %s\n", cmdio.HiBlack(ctx, strings.Repeat("─", col+10+autostopCol+12)))
 
 			for _, e := range entries {
 				id := e.SandboxID
 				def := ""
 				if id == defaultID {
-					def = accent("*")
+					def = cmdio.Cyan(ctx, "*")
 				}
-				// Pad ID manually to avoid ANSI codes breaking alignment.
+				// Pad ID manually so visible-width alignment is preserved
+				// after the helpers wrap each cell with ANSI escapes.
 				idPad := max(col-len(id), 0)
-				st := status(e.Status)
-				// Pad status to 10 visible chars.
+				st := status(ctx, e.Status)
 				stPad := max(10-len(e.Status), 0)
 				as := e.autoStopLabel()
 				asPad := max(autostopCol-len(as), 0)
-				idStr := bold(id)
-				if strings.EqualFold(e.Status, "running") {
-					idStr = cyan + bo + id + rs
-				}
+				idStr := cmdio.Cyan(ctx, id)
 				fmt.Fprintf(out, "  %s%s  %s%s  %s%s  %s\n",
 					idStr, strings.Repeat(" ", idPad),
 					st, strings.Repeat(" ", stPad),
-					dim(as), strings.Repeat(" ", asPad),
+					cmdio.HiBlack(ctx, as), strings.Repeat(" ", asPad),
 					def)
 			}
 			blank(out)
diff --git a/cmd/lakebox/register.go b/cmd/lakebox/register.go
index 8fdf42e2a7..0df618e767 100644
--- a/cmd/lakebox/register.go
+++ b/cmd/lakebox/register.go
@@ -10,6 +10,7 @@ import (
 
 	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdctx"
+	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/env"
 	"github.com/spf13/cobra"
 )
@@ -44,9 +45,9 @@ Example:
 
 			stderr := cmd.ErrOrStderr()
 			if generated {
-				ok(ctx, "Generated SSH key at "+dim(keyPath))
+				ok(ctx, "Generated SSH key at "+cmdio.HiBlack(ctx, keyPath))
 			} else {
-				field(stderr, "key", keyPath)
+				field(ctx, stderr, "key", keyPath)
 			}
 
 			pubKeyData, err := os.ReadFile(keyPath + ".pub")
@@ -62,7 +63,7 @@ Example:
 			s.ok("SSH key registered")
 
 			blank(stderr)
-			fmt.Fprintf(stderr, "  Run %s to connect.\n\n", bold("databricks lakebox ssh"))
+			fmt.Fprintf(stderr, "  Run %s to connect.\n\n", cmdio.Cyan(ctx, "databricks lakebox ssh"))
 			return nil
 		},
 	}
diff --git a/cmd/lakebox/ssh.go b/cmd/lakebox/ssh.go
index bed21de997..3f83b4b95e 100644
--- a/cmd/lakebox/ssh.go
+++ b/cmd/lakebox/ssh.go
@@ -11,6 +11,7 @@ import (
 
 	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdctx"
+	"github.com/databricks/cli/libs/cmdio"
 	"github.com/spf13/cobra"
 )
 
@@ -104,7 +105,7 @@ Examples:
 						return fmt.Errorf("failed to create lakebox: %w", err)
 					}
 					lakeboxID = result.SandboxID
-					s.ok("Lakebox " + bold(lakeboxID) + " ready")
+					s.ok("Lakebox " + cmdio.Cyan(ctx, lakeboxID) + " ready")
 
 					if err := setDefault(ctx, profile, lakeboxID); err != nil {
 						warn(ctx, fmt.Sprintf("Could not save default: %v", err))
@@ -117,8 +118,8 @@ Examples:
 				host = resolveGatewayHost(w.Config.Host)
 			}
 
-			s := spin(ctx, "Connecting to "+bold(lakeboxID)+"…")
-			s.ok("Connected to " + bold(lakeboxID))
+			s := spin(ctx, "Connecting to "+cmdio.Cyan(ctx, lakeboxID)+"…")
+			s.ok("Connected to " + cmdio.Cyan(ctx, lakeboxID))
 			return execSSHDirect(lakeboxID, host, gatewayPort, keyPath, extraArgs)
 		},
 	}
diff --git a/cmd/lakebox/status.go b/cmd/lakebox/status.go
index 7050c53ff6..50f4c037f1 100644
--- a/cmd/lakebox/status.go
+++ b/cmd/lakebox/status.go
@@ -6,6 +6,7 @@ import (
 
 	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdctx"
+	"github.com/databricks/cli/libs/cmdio"
 	"github.com/spf13/cobra"
 )
 
@@ -42,12 +43,12 @@ Example:
 
 			out := cmd.OutOrStdout()
 			blank(out)
-			field(out, "id", bold(entry.SandboxID))
-			field(out, "status", status(entry.Status))
+			field(ctx, out, "id", cmdio.Cyan(ctx, entry.SandboxID))
+			field(ctx, out, "status", status(ctx, entry.Status))
 			if entry.FQDN != "" {
-				field(out, "fqdn", dim(entry.FQDN))
+				field(ctx, out, "fqdn", cmdio.HiBlack(ctx, entry.FQDN))
 			}
-			field(out, "autostop", dim(entry.autoStopLabel()))
+			field(ctx, out, "autostop", cmdio.HiBlack(ctx, entry.autoStopLabel()))
 			blank(out)
 			return nil
 		},
diff --git a/cmd/lakebox/ui.go b/cmd/lakebox/ui.go
index 3d722986c3..066bb0a036 100644
--- a/cmd/lakebox/ui.go
+++ b/cmd/lakebox/ui.go
@@ -9,17 +9,6 @@ import (
 	"github.com/databricks/cli/libs/cmdio"
 )
 
-// ANSI escapes for inline highlighting. cmdio handles terminal capability
-// detection for the spinner, so we don't gate these on TTY here — strings
-// piped to a non-terminal still carry the codes, matching the behavior of
-// other CLI commands that call bold/dim helpers.
-const (
-	rs   = "\033[0m"  // reset
-	bo   = "\033[1m"  // bold
-	dm   = "\033[2m"  // dim
-	cyan = "\033[36m" // accent
-)
-
 // spinner wraps cmdio.NewSpinner with terminal ok/fail markers. After the
 // first call to ok or fail, the spinner is closed and a final line is logged
 // to stderr; subsequent calls are no-ops.
@@ -44,43 +33,39 @@ func (s *spinner) done(mark, msg string) {
 	}
 	s.finished = true
 	s.close()
-	cmdio.LogString(s.ctx, "  "+cyan+mark+rs+" "+msg)
+	cmdio.LogString(s.ctx, "  "+cmdio.Cyan(s.ctx, mark)+" "+msg)
 }
 
-// status formats a lakebox lifecycle status with the accent color.
-func status(s string) string {
+// status formats a lakebox lifecycle status with a color hint.
+func status(ctx context.Context, s string) string {
 	switch strings.ToLower(s) {
 	case "running":
-		return cyan + "running" + rs
+		return cmdio.Cyan(ctx, "running")
 	case "stopped":
-		return dm + "stopped" + rs
+		return cmdio.HiBlack(ctx, "stopped")
 	case "creating":
-		return cyan + bo + "creating…" + rs
+		return cmdio.Cyan(ctx, "creating…")
 	default:
-		return dm + strings.ToLower(s) + rs
+		return cmdio.HiBlack(ctx, strings.ToLower(s))
 	}
 }
 
-// field prints "  label  value" to w.
-func field(w io.Writer, label, value string) {
-	fmt.Fprintf(w, "  %s%-10s%s %s\n", dm, label, rs, value)
+// field prints "  label  value" to w, where label is dimmed.
+func field(ctx context.Context, w io.Writer, label, value string) {
+	fmt.Fprintf(w, "  %-10s %s\n", cmdio.HiBlack(ctx, label), value)
 }
 
 // ok prints "  ✓ message" to stderr via the cmdio context.
 func ok(ctx context.Context, msg string) {
-	cmdio.LogString(ctx, "  "+cyan+"✓"+rs+" "+msg)
+	cmdio.LogString(ctx, "  "+cmdio.Cyan(ctx, "✓")+" "+msg)
 }
 
 // warn prints "  ! message" to stderr via the cmdio context.
 func warn(ctx context.Context, msg string) {
-	cmdio.LogString(ctx, "  "+cyan+"!"+rs+" "+msg)
+	cmdio.LogString(ctx, "  "+cmdio.Cyan(ctx, "!")+" "+msg)
 }
 
 // blank prints an empty line to w.
 func blank(w io.Writer) {
 	fmt.Fprintln(w)
 }
-
-func accent(s string) string { return cyan + s + rs }
-func bold(s string) string   { return bo + s + rs }
-func dim(s string) string    { return dm + s + rs }

From 43807fadf3d58b05ee336c19e29a9b7f22bfb4a2 Mon Sep 17 00:00:00 2001
From: Pieter Noordhuis <pieter.noordhuis@databricks.com>
Date: Wed, 6 May 2026 15:57:39 +0200
Subject: [PATCH 18/36] cmdio: add Bold and Dim color helpers; restore lakebox
 parity
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

cmdio.Cyan/HiBlack covered most of lakebox's needs but conflated two
distinct visual roles: bold-for-emphasis (uncolored) on IDs and command
names, and dim (\x1b[2m, faint) on secondary metadata. The previous
commit collapsed both into Cyan/HiBlack and changed the rendering.

Add Bold and Dim helpers alongside the existing color set — ansiBold
already lived in color.go; ansiDim is new. Both gate on the same
SupportsStdoutColor capability check as Red/Green/etc.

Switch lakebox call sites back to Bold for IDs and command emphasis,
and to Dim for secondary text (autostop labels, FQDNs, "(default
cleared)", table headers, the "No lakeboxes found" notice). Running
lakebox IDs in `list` go back to bold-cyan via composition.

Co-authored-by: Isaac
---
 cmd/lakebox/config.go    |  4 ++--
 cmd/lakebox/create.go    |  2 +-
 cmd/lakebox/delete.go    |  4 ++--
 cmd/lakebox/list.go      | 13 ++++++++-----
 cmd/lakebox/register.go  |  4 ++--
 cmd/lakebox/ssh.go       |  6 +++---
 cmd/lakebox/status.go    |  6 +++---
 cmd/lakebox/ui.go        |  6 +++---
 libs/cmdio/color.go      |  7 +++++++
 libs/cmdio/color_test.go |  2 ++
 10 files changed, 33 insertions(+), 21 deletions(-)

diff --git a/cmd/lakebox/config.go b/cmd/lakebox/config.go
index f3da157e00..963e5a092c 100644
--- a/cmd/lakebox/config.go
+++ b/cmd/lakebox/config.go
@@ -88,8 +88,8 @@ Examples:
 			}
 
 			blank(out)
-			field(ctx, out, "id", cmdio.Cyan(ctx, updated.SandboxID))
-			field(ctx, out, "autostop", cmdio.HiBlack(ctx, updated.autoStopLabel()))
+			field(ctx, out, "id", cmdio.Bold(ctx, updated.SandboxID))
+			field(ctx, out, "autostop", cmdio.Dim(ctx, updated.autoStopLabel()))
 			blank(out)
 			return nil
 		},
diff --git a/cmd/lakebox/create.go b/cmd/lakebox/create.go
index 62bcb3085d..5303dabc30 100644
--- a/cmd/lakebox/create.go
+++ b/cmd/lakebox/create.go
@@ -46,7 +46,7 @@ Example:
 				return fmt.Errorf("failed to create lakebox: %w", err)
 			}
 
-			s.ok("Lakebox " + cmdio.Cyan(ctx, result.SandboxID) + " is " + status(ctx, result.Status))
+			s.ok("Lakebox " + cmdio.Bold(ctx, result.SandboxID) + " is " + status(ctx, result.Status))
 
 			profile := w.Config.Profile
 			if profile == "" {
diff --git a/cmd/lakebox/delete.go b/cmd/lakebox/delete.go
index 54d7a59e15..f589d3c986 100644
--- a/cmd/lakebox/delete.go
+++ b/cmd/lakebox/delete.go
@@ -40,9 +40,9 @@ Example:
 			}
 			if getDefault(ctx, profile) == lakeboxID {
 				_ = clearDefault(ctx, profile)
-				s.ok("Removed " + cmdio.Cyan(ctx, lakeboxID) + " " + cmdio.HiBlack(ctx, "(default cleared)"))
+				s.ok("Removed " + cmdio.Bold(ctx, lakeboxID) + " " + cmdio.Dim(ctx, "(default cleared)"))
 			} else {
-				s.ok("Removed " + cmdio.Cyan(ctx, lakeboxID))
+				s.ok("Removed " + cmdio.Bold(ctx, lakeboxID))
 			}
 			return nil
 		},
diff --git a/cmd/lakebox/list.go b/cmd/lakebox/list.go
index d3d03f77fe..d9c18e6d21 100644
--- a/cmd/lakebox/list.go
+++ b/cmd/lakebox/list.go
@@ -43,7 +43,7 @@ Example:
 			}
 
 			if len(entries) == 0 {
-				fmt.Fprintf(cmd.ErrOrStderr(), "  %s\n", cmdio.HiBlack(ctx, "No lakeboxes found."))
+				fmt.Fprintf(cmd.ErrOrStderr(), "  %s\n", cmdio.Dim(ctx, "No lakeboxes found."))
 				return nil
 			}
 
@@ -73,8 +73,8 @@ Example:
 			blank(out)
 			header := fmt.Sprintf("%-*s  %-10s  %-*s  %s",
 				col, "ID", "STATUS", autostopCol, "AUTOSTOP", "DEFAULT")
-			fmt.Fprintf(out, "  %s\n", cmdio.HiBlack(ctx, header))
-			fmt.Fprintf(out, "  %s\n", cmdio.HiBlack(ctx, strings.Repeat("─", col+10+autostopCol+12)))
+			fmt.Fprintf(out, "  %s\n", cmdio.Dim(ctx, header))
+			fmt.Fprintf(out, "  %s\n", cmdio.Dim(ctx, strings.Repeat("─", col+10+autostopCol+12)))
 
 			for _, e := range entries {
 				id := e.SandboxID
@@ -89,11 +89,14 @@ Example:
 				stPad := max(10-len(e.Status), 0)
 				as := e.autoStopLabel()
 				asPad := max(autostopCol-len(as), 0)
-				idStr := cmdio.Cyan(ctx, id)
+				idStr := cmdio.Bold(ctx, id)
+				if strings.EqualFold(e.Status, "running") {
+					idStr = cmdio.Bold(ctx, cmdio.Cyan(ctx, id))
+				}
 				fmt.Fprintf(out, "  %s%s  %s%s  %s%s  %s\n",
 					idStr, strings.Repeat(" ", idPad),
 					st, strings.Repeat(" ", stPad),
-					cmdio.HiBlack(ctx, as), strings.Repeat(" ", asPad),
+					cmdio.Dim(ctx, as), strings.Repeat(" ", asPad),
 					def)
 			}
 			blank(out)
diff --git a/cmd/lakebox/register.go b/cmd/lakebox/register.go
index 0df618e767..fbc09acd3b 100644
--- a/cmd/lakebox/register.go
+++ b/cmd/lakebox/register.go
@@ -45,7 +45,7 @@ Example:
 
 			stderr := cmd.ErrOrStderr()
 			if generated {
-				ok(ctx, "Generated SSH key at "+cmdio.HiBlack(ctx, keyPath))
+				ok(ctx, "Generated SSH key at "+cmdio.Dim(ctx, keyPath))
 			} else {
 				field(ctx, stderr, "key", keyPath)
 			}
@@ -63,7 +63,7 @@ Example:
 			s.ok("SSH key registered")
 
 			blank(stderr)
-			fmt.Fprintf(stderr, "  Run %s to connect.\n\n", cmdio.Cyan(ctx, "databricks lakebox ssh"))
+			fmt.Fprintf(stderr, "  Run %s to connect.\n\n", cmdio.Bold(ctx, "databricks lakebox ssh"))
 			return nil
 		},
 	}
diff --git a/cmd/lakebox/ssh.go b/cmd/lakebox/ssh.go
index 3f83b4b95e..5a52414123 100644
--- a/cmd/lakebox/ssh.go
+++ b/cmd/lakebox/ssh.go
@@ -105,7 +105,7 @@ Examples:
 						return fmt.Errorf("failed to create lakebox: %w", err)
 					}
 					lakeboxID = result.SandboxID
-					s.ok("Lakebox " + cmdio.Cyan(ctx, lakeboxID) + " ready")
+					s.ok("Lakebox " + cmdio.Bold(ctx, lakeboxID) + " ready")
 
 					if err := setDefault(ctx, profile, lakeboxID); err != nil {
 						warn(ctx, fmt.Sprintf("Could not save default: %v", err))
@@ -118,8 +118,8 @@ Examples:
 				host = resolveGatewayHost(w.Config.Host)
 			}
 
-			s := spin(ctx, "Connecting to "+cmdio.Cyan(ctx, lakeboxID)+"…")
-			s.ok("Connected to " + cmdio.Cyan(ctx, lakeboxID))
+			s := spin(ctx, "Connecting to "+cmdio.Bold(ctx, lakeboxID)+"…")
+			s.ok("Connected to " + cmdio.Bold(ctx, lakeboxID))
 			return execSSHDirect(lakeboxID, host, gatewayPort, keyPath, extraArgs)
 		},
 	}
diff --git a/cmd/lakebox/status.go b/cmd/lakebox/status.go
index 50f4c037f1..1e428d79ff 100644
--- a/cmd/lakebox/status.go
+++ b/cmd/lakebox/status.go
@@ -43,12 +43,12 @@ Example:
 
 			out := cmd.OutOrStdout()
 			blank(out)
-			field(ctx, out, "id", cmdio.Cyan(ctx, entry.SandboxID))
+			field(ctx, out, "id", cmdio.Bold(ctx, entry.SandboxID))
 			field(ctx, out, "status", status(ctx, entry.Status))
 			if entry.FQDN != "" {
-				field(ctx, out, "fqdn", cmdio.HiBlack(ctx, entry.FQDN))
+				field(ctx, out, "fqdn", cmdio.Dim(ctx, entry.FQDN))
 			}
-			field(ctx, out, "autostop", cmdio.HiBlack(ctx, entry.autoStopLabel()))
+			field(ctx, out, "autostop", cmdio.Dim(ctx, entry.autoStopLabel()))
 			blank(out)
 			return nil
 		},
diff --git a/cmd/lakebox/ui.go b/cmd/lakebox/ui.go
index 066bb0a036..709d72586f 100644
--- a/cmd/lakebox/ui.go
+++ b/cmd/lakebox/ui.go
@@ -42,17 +42,17 @@ func status(ctx context.Context, s string) string {
 	case "running":
 		return cmdio.Cyan(ctx, "running")
 	case "stopped":
-		return cmdio.HiBlack(ctx, "stopped")
+		return cmdio.Dim(ctx, "stopped")
 	case "creating":
 		return cmdio.Cyan(ctx, "creating…")
 	default:
-		return cmdio.HiBlack(ctx, strings.ToLower(s))
+		return cmdio.Dim(ctx, strings.ToLower(s))
 	}
 }
 
 // field prints "  label  value" to w, where label is dimmed.
 func field(ctx context.Context, w io.Writer, label, value string) {
-	fmt.Fprintf(w, "  %-10s %s\n", cmdio.HiBlack(ctx, label), value)
+	fmt.Fprintf(w, "  %-10s %s\n", cmdio.Dim(ctx, label), value)
 }
 
 // ok prints "  ✓ message" to stderr via the cmdio context.
diff --git a/libs/cmdio/color.go b/libs/cmdio/color.go
index 4066b30f75..a2a7ce24e2 100644
--- a/libs/cmdio/color.go
+++ b/libs/cmdio/color.go
@@ -11,6 +11,7 @@ import (
 const (
 	ansiReset     = "\x1b[0m"
 	ansiBold      = "\x1b[1m"
+	ansiDim       = "\x1b[2m"
 	ansiItalic    = "\x1b[3m"
 	ansiRed       = "\x1b[31m"
 	ansiGreen     = "\x1b[32m"
@@ -42,6 +43,12 @@ func render(ctx context.Context, code, msg string) string {
 	return code + msg + ansiReset
 }
 
+// Bold renders msg in bold.
+func Bold(ctx context.Context, msg string) string { return render(ctx, ansiBold, msg) }
+
+// Dim renders msg in dim (faint) intensity.
+func Dim(ctx context.Context, msg string) string { return render(ctx, ansiDim, msg) }
+
 // Red renders msg in red.
 func Red(ctx context.Context, msg string) string { return render(ctx, ansiRed, msg) }
 
diff --git a/libs/cmdio/color_test.go b/libs/cmdio/color_test.go
index 54df185982..dcc45f9c94 100644
--- a/libs/cmdio/color_test.go
+++ b/libs/cmdio/color_test.go
@@ -27,6 +27,8 @@ func TestColorHelpersEmitSGRWhenEnabled(t *testing.T) {
 		got  string
 		want string
 	}{
+		{"Bold", cmdio.Bold(ctx, "id"), "\x1b[1mid\x1b[0m"},
+		{"Dim", cmdio.Dim(ctx, "hint"), "\x1b[2mhint\x1b[0m"},
 		{"Red", cmdio.Red(ctx, "hello"), "\x1b[31mhello\x1b[0m"},
 		{"Green", cmdio.Green(ctx, "ok"), "\x1b[32mok\x1b[0m"},
 		{"Yellow", cmdio.Yellow(ctx, "warn"), "\x1b[33mwarn\x1b[0m"},

From 205edce567a02772074e2b9d37bf1fe37c5b8234 Mon Sep 17 00:00:00 2001
From: Pieter Noordhuis <pieter.noordhuis@databricks.com>
Date: Wed, 6 May 2026 16:59:03 +0200
Subject: [PATCH 19/36] lakebox: restore status('creating') bold and field
 column alignment
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Two parity gaps from the previous commit, found by diffing byte-level
output against the original ui.go:

- status('creating') was bold cyan in the original; the cmdio rewrite
  dropped the bold. Restore via Bold(Cyan(...)) composition. Bytes
  differ from the original ('\x1b[1m\x1b[36m...\x1b[0m\x1b[0m' vs
  '\x1b[36m\x1b[1m...\x1b[0m') but render identically — SGR codes
  are additive, the extra trailing reset is a no-op.

- field() applied %-10s padding to the already-Dim-wrapped label, so
  the SGR escapes inflated the byte count and column alignment broke
  whenever color was enabled. Pad first, then wrap.

Co-authored-by: Isaac
---
 cmd/lakebox/ui.go | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/cmd/lakebox/ui.go b/cmd/lakebox/ui.go
index 709d72586f..b9ce334cd8 100644
--- a/cmd/lakebox/ui.go
+++ b/cmd/lakebox/ui.go
@@ -44,15 +44,17 @@ func status(ctx context.Context, s string) string {
 	case "stopped":
 		return cmdio.Dim(ctx, "stopped")
 	case "creating":
-		return cmdio.Cyan(ctx, "creating…")
+		return cmdio.Bold(ctx, cmdio.Cyan(ctx, "creating…"))
 	default:
 		return cmdio.Dim(ctx, strings.ToLower(s))
 	}
 }
 
-// field prints "  label  value" to w, where label is dimmed.
+// field prints "  label  value" to w, where label is dimmed and padded to a
+// fixed visible width. Padding has to happen before Dim so the SGR escapes
+// don't inflate the byte count and break column alignment.
 func field(ctx context.Context, w io.Writer, label, value string) {
-	fmt.Fprintf(w, "  %-10s %s\n", cmdio.Dim(ctx, label), value)
+	fmt.Fprintf(w, "  %s %s\n", cmdio.Dim(ctx, fmt.Sprintf("%-10s", label)), value)
 }
 
 // ok prints "  ✓ message" to stderr via the cmdio context.

From f66fe2ab28107b52597ffdd63861d36266e4592d Mon Sep 17 00:00:00 2001
From: Pieter Noordhuis <pieter.noordhuis@databricks.com>
Date: Thu, 7 May 2026 10:27:07 +0200
Subject: [PATCH 20/36] lakebox: drop unix-only exec_unix.go and runtime.GOOS
 branch
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The Unix path used syscall.Exec to replace the Go process with ssh
directly, saving one fork. The Windows path already used
exec.Command(...).Run(), and that works on all platforms — terminal
signals are delivered to ssh via the foreground process group either
way. Collapse to one cross-platform path; drop the build-tagged file
and the runtime.GOOS check.

Co-authored-by: Isaac
---
 cmd/lakebox/exec_unix.go | 13 -------------
 cmd/lakebox/ssh.go       | 26 ++++++++------------------
 2 files changed, 8 insertions(+), 31 deletions(-)
 delete mode 100644 cmd/lakebox/exec_unix.go

diff --git a/cmd/lakebox/exec_unix.go b/cmd/lakebox/exec_unix.go
deleted file mode 100644
index d47f629572..0000000000
--- a/cmd/lakebox/exec_unix.go
+++ /dev/null
@@ -1,13 +0,0 @@
-//go:build !windows
-
-package lakebox
-
-import (
-	"os"
-	"syscall"
-)
-
-// execSyscall replaces the current process with the given command (Unix only).
-func execSyscall(path string, args []string) error {
-	return syscall.Exec(path, args, os.Environ())
-}
diff --git a/cmd/lakebox/ssh.go b/cmd/lakebox/ssh.go
index 5a52414123..b448516858 100644
--- a/cmd/lakebox/ssh.go
+++ b/cmd/lakebox/ssh.go
@@ -6,7 +6,6 @@ import (
 	"io/fs"
 	"os"
 	"os/exec"
-	"runtime"
 	"strings"
 
 	"github.com/databricks/cli/cmd/root"
@@ -130,16 +129,11 @@ Examples:
 	return cmd
 }
 
-// execSSHDirect execs into ssh with all options passed as args (no ~/.ssh/config needed).
-// Extra args are appended after the destination (for remote commands or ssh flags).
+// execSSHDirect runs ssh with all options passed as args (no ~/.ssh/config
+// needed). Extra args are appended after the destination for remote commands
+// or ssh flags.
 func execSSHDirect(lakeboxID, host, port, keyPath string, extraArgs []string) error {
-	sshPath, err := exec.LookPath("ssh")
-	if err != nil {
-		return fmt.Errorf("ssh not found in PATH: %w", err)
-	}
-
 	args := []string{
-		"ssh",
 		"-i", keyPath,
 		"-p", port,
 		"-o", "IdentitiesOnly=yes",
@@ -151,13 +145,9 @@ func execSSHDirect(lakeboxID, host, port, keyPath string, extraArgs []string) er
 	}
 	args = append(args, extraArgs...)
 
-	if runtime.GOOS == "windows" {
-		cmd := exec.Command(sshPath, args[1:]...)
-		cmd.Stdin = os.Stdin
-		cmd.Stdout = os.Stdout
-		cmd.Stderr = os.Stderr
-		return cmd.Run()
-	}
-
-	return execSyscall(sshPath, args)
+	cmd := exec.Command("ssh", args...)
+	cmd.Stdin = os.Stdin
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+	return cmd.Run()
 }

From 102f279754f2f0c16ffad03fed6450587b8db4f1 Mon Sep 17 00:00:00 2001
From: Pieter Noordhuis <pieter.noordhuis@databricks.com>
Date: Thu, 7 May 2026 10:29:10 +0200
Subject: [PATCH 21/36] lakebox: use libs/execv for ssh process replacement
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

libs/execv already wraps the syscall.Exec / Windows-emulation pattern
the previous version reimplemented inline. Switch to it so ssh truly
replaces the CLI process on Unix instead of running as a child — fewer
moving parts when the user hits Ctrl-C, and one fewer Go process in the
ps tree for the lifetime of the session.

Co-authored-by: Isaac
---
 cmd/lakebox/ssh.go | 19 ++++++++++---------
 1 file changed, 10 insertions(+), 9 deletions(-)

diff --git a/cmd/lakebox/ssh.go b/cmd/lakebox/ssh.go
index b448516858..1cc234646e 100644
--- a/cmd/lakebox/ssh.go
+++ b/cmd/lakebox/ssh.go
@@ -5,12 +5,12 @@ import (
 	"fmt"
 	"io/fs"
 	"os"
-	"os/exec"
 	"strings"
 
 	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/databricks/cli/libs/cmdio"
+	"github.com/databricks/cli/libs/execv"
 	"github.com/spf13/cobra"
 )
 
@@ -129,11 +129,13 @@ Examples:
 	return cmd
 }
 
-// execSSHDirect runs ssh with all options passed as args (no ~/.ssh/config
-// needed). Extra args are appended after the destination for remote commands
-// or ssh flags.
+// execSSHDirect replaces the CLI process with ssh (or simulates that on
+// Windows via execv). All options are passed on the command line, so no
+// ~/.ssh/config entry is required. Extra args are appended after the
+// destination for remote commands or ssh flags.
 func execSSHDirect(lakeboxID, host, port, keyPath string, extraArgs []string) error {
 	args := []string{
+		"ssh",
 		"-i", keyPath,
 		"-p", port,
 		"-o", "IdentitiesOnly=yes",
@@ -145,9 +147,8 @@ func execSSHDirect(lakeboxID, host, port, keyPath string, extraArgs []string) er
 	}
 	args = append(args, extraArgs...)
 
-	cmd := exec.Command("ssh", args...)
-	cmd.Stdin = os.Stdin
-	cmd.Stdout = os.Stdout
-	cmd.Stderr = os.Stderr
-	return cmd.Run()
+	return execv.Execv(execv.Options{
+		Args: args,
+		Env:  os.Environ(),
+	})
 }

From 08a56bef2814ab70fbd2429b5b378cb4206cdc00 Mon Sep 17 00:00:00 2001
From: Pieter Noordhuis <pieter.noordhuis@databricks.com>
Date: Thu, 7 May 2026 10:55:03 +0200
Subject: [PATCH 22/36] lakebox: rewrite api.go on top of the SDK ApiClient

Replace the hand-rolled HTTP plumbing with client.DatabricksClient.Do,
following the pattern in cmd/api/api.go and bundle/deploy/filer.go. Each
method becomes a single Do() call; the SDK handles auth, JSON marshal,
JSON unmarshal, error parsing, and retries.

Removed:
 - doRequest (manual http.NewRequestWithContext + Config.Authenticate)
 - parseAPIError + the local apiError type (SDK returns apierr.APIError)
 - manual json.Marshal / json.NewDecoder.Decode in every method
 - net/http response status-code branching

Preserved:
 - X-Databricks-Org-Id is still injected on every call. The SDK's
   Config.WorkspaceID is the source of truth; we fall back to parsing
   `?o=<id>` off the host because some staging gateways are configured
   that way and the SDK doesn't lift the query into Config.WorkspaceID.

newLakeboxAPI now returns (*lakeboxAPI, error) since client.New can fail
on bad config; callers updated.

Co-authored-by: Isaac
---
 cmd/lakebox/api.go      | 256 ++++++++++++----------------------------
 cmd/lakebox/config.go   |   5 +-
 cmd/lakebox/create.go   |   6 +-
 cmd/lakebox/delete.go   |   6 +-
 cmd/lakebox/list.go     |   5 +-
 cmd/lakebox/register.go |   6 +-
 cmd/lakebox/status.go   |   5 +-
 7 files changed, 101 insertions(+), 188 deletions(-)

diff --git a/cmd/lakebox/api.go b/cmd/lakebox/api.go
index acaeff47e8..2b65fd8c70 100644
--- a/cmd/lakebox/api.go
+++ b/cmd/lakebox/api.go
@@ -1,26 +1,35 @@
 package lakebox
 
 import (
-	"bytes"
 	"context"
-	"encoding/json"
 	"fmt"
-	"io"
 	"net/http"
 	"net/url"
 	"strings"
 	"time"
 
 	"github.com/databricks/databricks-sdk-go"
+	"github.com/databricks/databricks-sdk-go/client"
 )
 
 // Sandboxes live under the `/sandboxes` sub-collection of the lakebox service
 // namespace (see `lakebox.proto` `LakeboxService.CreateSandbox`).
 const lakeboxAPIPath = "/api/2.0/lakebox/sandboxes"
 
-// lakeboxAPI wraps raw HTTP calls to the lakebox REST API.
+// SSH keys are nested under the lakebox service namespace alongside
+// `sandboxes/` (see `LakeboxService.CreateSshKey`).
+const lakeboxKeysAPIPath = "/api/2.0/lakebox/ssh-keys"
+
+// orgIDHeader is sent by multi-workspace gateways (e.g. dogfood staging) so
+// the gateway can scope the credential to a specific workspace. Without it,
+// requests fail with "Credential was not sent or was of an unsupported type
+// for this API."
+const orgIDHeader = "X-Databricks-Org-Id"
+
+// lakeboxAPI wraps the SDK ApiClient with workspace-id-aware request headers.
 type lakeboxAPI struct {
-	w *databricks.WorkspaceClient
+	c    *client.DatabricksClient
+	wsID string
 }
 
 // createRequest is the JSON body for POST /api/2.0/lakebox/sandboxes.
@@ -126,110 +135,94 @@ type listResponse struct {
 	Sandboxes []sandboxEntry `json:"sandboxes"`
 }
 
-// apiError is the error body returned by the lakebox API.
-type apiError struct {
-	ErrorCode string `json:"error_code"`
-	Message   string `json:"message"`
-}
-
-func (e *apiError) Error() string {
-	return fmt.Sprintf("%s: %s", e.ErrorCode, e.Message)
+// updateBody is the PATCH request body. The proto declares
+// `UpdateSandboxRequest { Sandbox sandbox = 1 }` with `body: "sandbox"`
+// in the (google.api.http) annotation, so the HTTP body is the inner
+// `Sandbox` message directly — there is no `{"sandbox": {...}}`
+// wrapping on the wire.
+//
+// Pointer fields encode the proto3 `optional` semantics — only the
+// fields we explicitly set are emitted, leaving everything else
+// server-untouched. `IdleTimeout` is a proto3-canonical Duration
+// string (e.g. `"900s"`); the server-side wire type is
+// `google.protobuf.Duration`.
+type updateBody struct {
+	SandboxID   string  `json:"sandbox_id"`
+	IdleTimeout *string `json:"idle_timeout,omitempty"`
+	NoAutostop  *bool   `json:"no_autostop,omitempty"`
 }
 
-func newLakeboxAPI(w *databricks.WorkspaceClient) *lakeboxAPI {
-	return &lakeboxAPI{w: w}
+// registerKeyRequest is the JSON body for POST /api/2.0/lakebox/ssh-keys.
+type registerKeyRequest struct {
+	PublicKey string `json:"public_key"`
+	Name      string `json:"name,omitempty"`
 }
 
-// create calls POST /api/2.0/lakebox with an optional public key.
-func (a *lakeboxAPI) create(ctx context.Context, publicKey string) (*createResponse, error) {
-	body := createRequest{PublicKey: publicKey}
-	jsonBody, err := json.Marshal(body)
+func newLakeboxAPI(w *databricks.WorkspaceClient) (*lakeboxAPI, error) {
+	c, err := client.New(w.Config)
 	if err != nil {
-		return nil, fmt.Errorf("failed to marshal request: %w", err)
+		return nil, fmt.Errorf("failed to create lakebox API client: %w", err)
 	}
+	return &lakeboxAPI{c: c, wsID: resolveWorkspaceID(w)}, nil
+}
 
-	resp, err := a.doRequest(ctx, "POST", lakeboxAPIPath, bytes.NewReader(jsonBody))
+// resolveWorkspaceID returns the workspace ID for the org-id header. Falls
+// back to the `?o=<id>` query parameter on the Host because some staging
+// gateways are configured that way and the SDK does not strip it into
+// Config.WorkspaceID.
+func resolveWorkspaceID(w *databricks.WorkspaceClient) string {
+	if id := w.Config.WorkspaceID; id != "" {
+		return id
+	}
+	parsed, err := url.Parse(w.Config.Host)
 	if err != nil {
-		return nil, err
+		return ""
 	}
-	defer resp.Body.Close()
+	return parsed.Query().Get("o")
+}
 
-	if resp.StatusCode != http.StatusOK {
-		return nil, parseAPIError(resp)
+func (a *lakeboxAPI) headers() map[string]string {
+	if a.wsID == "" {
+		return nil
 	}
+	return map[string]string{orgIDHeader: a.wsID}
+}
 
-	var result createResponse
-	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
-		return nil, fmt.Errorf("failed to decode response: %w", err)
+// create calls POST /api/2.0/lakebox/sandboxes with an optional public key.
+func (a *lakeboxAPI) create(ctx context.Context, publicKey string) (*createResponse, error) {
+	var resp createResponse
+	err := a.c.Do(ctx, http.MethodPost, lakeboxAPIPath, a.headers(), nil, createRequest{PublicKey: publicKey}, &resp)
+	if err != nil {
+		return nil, err
 	}
-	return &result, nil
+	return &resp, nil
 }
 
 // list calls GET /api/2.0/lakebox/sandboxes.
 func (a *lakeboxAPI) list(ctx context.Context) ([]sandboxEntry, error) {
-	resp, err := a.doRequest(ctx, "GET", lakeboxAPIPath, nil)
+	var resp listResponse
+	err := a.c.Do(ctx, http.MethodGet, lakeboxAPIPath, a.headers(), nil, nil, &resp)
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
-
-	if resp.StatusCode != http.StatusOK {
-		return nil, parseAPIError(resp)
-	}
-
-	var result listResponse
-	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
-		return nil, fmt.Errorf("failed to decode response: %w", err)
-	}
-	return result.Sandboxes, nil
+	return resp.Sandboxes, nil
 }
 
 // get calls GET /api/2.0/lakebox/sandboxes/{id}.
 func (a *lakeboxAPI) get(ctx context.Context, id string) (*sandboxEntry, error) {
-	resp, err := a.doRequest(ctx, "GET", lakeboxAPIPath+"/"+id, nil)
+	var resp sandboxEntry
+	err := a.c.Do(ctx, http.MethodGet, lakeboxAPIPath+"/"+id, a.headers(), nil, nil, &resp)
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
-
-	if resp.StatusCode != http.StatusOK {
-		return nil, parseAPIError(resp)
-	}
-
-	var result sandboxEntry
-	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
-		return nil, fmt.Errorf("failed to decode response: %w", err)
-	}
-	return &result, nil
-}
-
-// updateBody is the PATCH request body. The proto declares
-// `UpdateSandboxRequest { Sandbox sandbox = 1 }` with `body: "sandbox"`
-// in the (google.api.http) annotation, so the HTTP body is the inner
-// `Sandbox` message directly — there is no `{"sandbox": {...}}`
-// wrapping on the wire.
-//
-// Pointer fields encode the proto3 `optional` semantics — only the
-// fields we explicitly set are emitted, leaving everything else
-// server-untouched. `IdleTimeout` is a proto3-canonical Duration
-// string (e.g. `"900s"`); the server-side wire type is
-// `google.protobuf.Duration`.
-type updateBody struct {
-	SandboxID   string  `json:"sandbox_id"`
-	IdleTimeout *string `json:"idle_timeout,omitempty"`
-	NoAutostop  *bool   `json:"no_autostop,omitempty"`
+	return &resp, nil
 }
 
 // update calls PATCH /api/2.0/lakebox/sandboxes/{id} with whichever of
 // `idle_timeout` / `no_autostop` the caller chose to set. Fields left
 // nil are omitted from the wire payload, so the server preserves their
 // current values. Returns the refreshed `sandboxEntry`.
-func (a *lakeboxAPI) update(
-	ctx context.Context,
-	id string,
-	idleTimeoutSecs *int64,
-	noAutostop *bool,
-) (*sandboxEntry, error) {
+func (a *lakeboxAPI) update(ctx context.Context, id string, idleTimeoutSecs *int64, noAutostop *bool) (*sandboxEntry, error) {
 	var idleTimeout *string
 	if idleTimeoutSecs != nil {
 		s := fmt.Sprintf("%ds", *idleTimeoutSecs)
@@ -240,121 +233,20 @@ func (a *lakeboxAPI) update(
 		IdleTimeout: idleTimeout,
 		NoAutostop:  noAutostop,
 	}
-	jsonBody, err := json.Marshal(body)
-	if err != nil {
-		return nil, fmt.Errorf("failed to marshal request: %w", err)
-	}
-
-	resp, err := a.doRequest(ctx, "PATCH", lakeboxAPIPath+"/"+id, bytes.NewReader(jsonBody))
+	var resp sandboxEntry
+	err := a.c.Do(ctx, http.MethodPatch, lakeboxAPIPath+"/"+id, a.headers(), nil, body, &resp)
 	if err != nil {
 		return nil, err
 	}
-	defer resp.Body.Close()
-
-	if resp.StatusCode != http.StatusOK {
-		return nil, parseAPIError(resp)
-	}
-
-	var result sandboxEntry
-	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
-		return nil, fmt.Errorf("failed to decode response: %w", err)
-	}
-	return &result, nil
+	return &resp, nil
 }
 
 // delete calls DELETE /api/2.0/lakebox/sandboxes/{id}.
 func (a *lakeboxAPI) delete(ctx context.Context, id string) error {
-	resp, err := a.doRequest(ctx, "DELETE", lakeboxAPIPath+"/"+id, nil)
-	if err != nil {
-		return err
-	}
-	defer resp.Body.Close()
-
-	if resp.StatusCode != http.StatusOK {
-		return parseAPIError(resp)
-	}
-	return nil
-}
-
-// doRequest makes an authenticated HTTP request to the workspace.
-func (a *lakeboxAPI) doRequest(ctx context.Context, method, path string, body io.Reader) (*http.Response, error) {
-	// The configured host may be just a hostname or may carry a workspace
-	// selector in the query (e.g. `https://dogfood.staging.databricks.com/?o=...`).
-	// Parse it so we can append the API path while preserving the query, and so
-	// we can pull the workspace ID out of `?o=<id>` when the SDK config doesn't
-	// carry it on a separate `workspace_id` field.
-	parsed, err := url.Parse(a.w.Config.Host)
-	if err != nil {
-		return nil, fmt.Errorf("failed to parse host %q: %w", a.w.Config.Host, err)
-	}
-	wsid := a.w.Config.WorkspaceID
-	if wsid == "" {
-		if v := parsed.Query().Get("o"); v != "" {
-			wsid = v
-		}
-	}
-	parsed.Path = strings.TrimRight(parsed.Path, "/") + path
-
-	req, err := http.NewRequestWithContext(ctx, method, parsed.String(), body)
-	if err != nil {
-		return nil, fmt.Errorf("failed to create request: %w", err)
-	}
-
-	if err := a.w.Config.Authenticate(req); err != nil {
-		return nil, fmt.Errorf("failed to authenticate: %w", err)
-	}
-
-	// Multi-workspace gateways (e.g. dogfood.staging.databricks.com) need a
-	// workspace selector to route the request — without it the gateway can't
-	// scope the credential and rejects with "Credential was not sent or was of
-	// an unsupported type for this API". `?o=<id>` in the URL works as a
-	// fallback, but the explicit header is the well-defined contract.
-	if wsid != "" {
-		req.Header.Set("X-Databricks-Org-Id", wsid)
-	}
-
-	if body != nil {
-		req.Header.Set("Content-Type", "application/json")
-	}
-
-	return http.DefaultClient.Do(req)
-}
-
-func parseAPIError(resp *http.Response) error {
-	body, _ := io.ReadAll(resp.Body)
-	var apiErr apiError
-	if json.Unmarshal(body, &apiErr) == nil && apiErr.Message != "" {
-		return &apiErr
-	}
-	return fmt.Errorf("HTTP %d: %s", resp.StatusCode, string(body))
-}
-
-// SSH keys are now nested under the lakebox service namespace alongside
-// `sandboxes/` (see `LakeboxService.CreateSshKey`).
-const lakeboxKeysAPIPath = "/api/2.0/lakebox/ssh-keys"
-
-// registerKeyRequest is the JSON body for POST /api/2.0/lakebox/ssh-keys.
-type registerKeyRequest struct {
-	PublicKey string `json:"public_key"`
-	Name      string `json:"name,omitempty"`
+	return a.c.Do(ctx, http.MethodDelete, lakeboxAPIPath+"/"+id, a.headers(), nil, nil, nil)
 }
 
 // registerKey calls POST /api/2.0/lakebox/ssh-keys.
 func (a *lakeboxAPI) registerKey(ctx context.Context, publicKey string) error {
-	body := registerKeyRequest{PublicKey: publicKey}
-	jsonBody, err := json.Marshal(body)
-	if err != nil {
-		return fmt.Errorf("failed to marshal request: %w", err)
-	}
-
-	resp, err := a.doRequest(ctx, "POST", lakeboxKeysAPIPath, bytes.NewReader(jsonBody))
-	if err != nil {
-		return err
-	}
-	defer resp.Body.Close()
-
-	if resp.StatusCode != http.StatusOK {
-		return parseAPIError(resp)
-	}
-	return nil
+	return a.c.Do(ctx, http.MethodPost, lakeboxKeysAPIPath, a.headers(), nil, registerKeyRequest{PublicKey: publicKey}, nil)
 }
diff --git a/cmd/lakebox/config.go b/cmd/lakebox/config.go
index 963e5a092c..2861930cc6 100644
--- a/cmd/lakebox/config.go
+++ b/cmd/lakebox/config.go
@@ -56,7 +56,10 @@ Examples:
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := cmd.Context()
 			w := cmdctx.WorkspaceClient(ctx)
-			api := newLakeboxAPI(w)
+			api, err := newLakeboxAPI(w)
+			if err != nil {
+				return err
+			}
 			out := cmd.OutOrStdout()
 
 			id := args[0]
diff --git a/cmd/lakebox/create.go b/cmd/lakebox/create.go
index 5303dabc30..ea5c47cac6 100644
--- a/cmd/lakebox/create.go
+++ b/cmd/lakebox/create.go
@@ -27,7 +27,10 @@ Example:
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := cmd.Context()
 			w := cmdctx.WorkspaceClient(ctx)
-			api := newLakeboxAPI(w)
+			api, err := newLakeboxAPI(w)
+			if err != nil {
+				return err
+			}
 
 			var publicKey string
 			if publicKeyFile != "" {
@@ -39,6 +42,7 @@ Example:
 			}
 
 			s := spin(ctx, "Provisioning your lakebox…")
+			defer s.Close()
 
 			result, err := api.create(ctx, publicKey)
 			if err != nil {
diff --git a/cmd/lakebox/delete.go b/cmd/lakebox/delete.go
index f589d3c986..001a382252 100644
--- a/cmd/lakebox/delete.go
+++ b/cmd/lakebox/delete.go
@@ -24,10 +24,14 @@ Example:
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := cmd.Context()
 			w := cmdctx.WorkspaceClient(ctx)
-			api := newLakeboxAPI(w)
+			api, err := newLakeboxAPI(w)
+			if err != nil {
+				return err
+			}
 
 			lakeboxID := args[0]
 			s := spin(ctx, "Removing "+lakeboxID+"…")
+			defer s.Close()
 
 			if err := api.delete(ctx, lakeboxID); err != nil {
 				s.fail("Failed to delete " + lakeboxID)
diff --git a/cmd/lakebox/list.go b/cmd/lakebox/list.go
index d9c18e6d21..6dc1b42fb1 100644
--- a/cmd/lakebox/list.go
+++ b/cmd/lakebox/list.go
@@ -29,7 +29,10 @@ Example:
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := cmd.Context()
 			w := cmdctx.WorkspaceClient(ctx)
-			api := newLakeboxAPI(w)
+			api, err := newLakeboxAPI(w)
+			if err != nil {
+				return err
+			}
 
 			entries, err := api.list(ctx)
 			if err != nil {
diff --git a/cmd/lakebox/register.go b/cmd/lakebox/register.go
index fbc09acd3b..c3e34e4ea3 100644
--- a/cmd/lakebox/register.go
+++ b/cmd/lakebox/register.go
@@ -36,7 +36,10 @@ Example:
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := cmd.Context()
 			w := cmdctx.WorkspaceClient(ctx)
-			api := newLakeboxAPI(w)
+			api, err := newLakeboxAPI(w)
+			if err != nil {
+				return err
+			}
 
 			keyPath, generated, err := ensureLakeboxKey(ctx)
 			if err != nil {
@@ -56,6 +59,7 @@ Example:
 			}
 
 			s := spin(ctx, "Registering key…")
+			defer s.Close()
 			if err := api.registerKey(ctx, string(pubKeyData)); err != nil {
 				s.fail("Failed to register key")
 				return fmt.Errorf("failed to register key: %w", err)
diff --git a/cmd/lakebox/status.go b/cmd/lakebox/status.go
index 1e428d79ff..ee9c276aa8 100644
--- a/cmd/lakebox/status.go
+++ b/cmd/lakebox/status.go
@@ -26,7 +26,10 @@ Example:
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := cmd.Context()
 			w := cmdctx.WorkspaceClient(ctx)
-			api := newLakeboxAPI(w)
+			api, err := newLakeboxAPI(w)
+			if err != nil {
+				return err
+			}
 
 			lakeboxID := args[0]
 

From 205d54597b37e5a7dfdd95f799c7ffecbee1fae9 Mon Sep 17 00:00:00 2001
From: Pieter Noordhuis <pieter.noordhuis@databricks.com>
Date: Thu, 7 May 2026 10:55:12 +0200
Subject: [PATCH 23/36] lakebox: make spinner Close() idempotent; defer it at
 every spin site
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Today if a code path between spin(...) and s.ok/s.fail returns early,
the cmdio Bubble Tea program keeps running and we leak a goroutine
plus garble the terminal. The wrapper kept its own `finished` gate but
exposed no way to close without printing a marker.

Add Close() that stops the spinner with no marker (wired through the
same `finished` gate, so calling Close() after ok/fail is a no-op),
and `defer s.Close()` at every spin site. ok/fail still print the ✓/✗
line on the success/failure paths; Close is just the cleanup safety
net.

Co-authored-by: Isaac
---
 cmd/lakebox/ssh.go |  7 ++++++-
 cmd/lakebox/ui.go  | 24 +++++++++++++++++-------
 2 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/cmd/lakebox/ssh.go b/cmd/lakebox/ssh.go
index 1cc234646e..2bbd5a7e34 100644
--- a/cmd/lakebox/ssh.go
+++ b/cmd/lakebox/ssh.go
@@ -91,13 +91,17 @@ Examples:
 				if def := getDefault(ctx, profile); def != "" {
 					lakeboxID = def
 				} else {
-					api := newLakeboxAPI(w)
+					api, err := newLakeboxAPI(w)
+					if err != nil {
+						return err
+					}
 					pubKeyData, err := os.ReadFile(keyPath + ".pub")
 					if err != nil {
 						return fmt.Errorf("failed to read public key %s.pub: %w", keyPath, err)
 					}
 
 					s := spin(ctx, "Provisioning your lakebox…")
+					defer s.Close()
 					result, err := api.create(ctx, string(pubKeyData))
 					if err != nil {
 						s.fail("Failed to create lakebox")
@@ -118,6 +122,7 @@ Examples:
 			}
 
 			s := spin(ctx, "Connecting to "+cmdio.Bold(ctx, lakeboxID)+"…")
+			defer s.Close()
 			s.ok("Connected to " + cmdio.Bold(ctx, lakeboxID))
 			return execSSHDirect(lakeboxID, host, gatewayPort, keyPath, extraArgs)
 		},
diff --git a/cmd/lakebox/ui.go b/cmd/lakebox/ui.go
index b9ce334cd8..c2669ce8ef 100644
--- a/cmd/lakebox/ui.go
+++ b/cmd/lakebox/ui.go
@@ -9,31 +9,41 @@ import (
 	"github.com/databricks/cli/libs/cmdio"
 )
 
-// spinner wraps cmdio.NewSpinner with terminal ok/fail markers. After the
-// first call to ok or fail, the spinner is closed and a final line is logged
-// to stderr; subsequent calls are no-ops.
+// spinner wraps cmdio.NewSpinner with terminal ok/fail markers. The first
+// call to ok, fail, or Close closes the underlying cmdio spinner; ok/fail
+// also log a final line to stderr. Subsequent calls are no-ops, so callers
+// are expected to `defer s.Close()` and call ok/fail on the success/failure
+// path. Close on its own (no marker) just stops the spinner — useful when an
+// error path returns before reaching ok/fail.
 type spinner struct {
 	ctx      context.Context
-	close    func()
+	inner    func()
 	finished bool
 }
 
 func spin(ctx context.Context, msg string) *spinner {
 	sp := cmdio.NewSpinner(ctx)
 	sp.Update(msg)
-	return &spinner{ctx: ctx, close: sp.Close}
+	return &spinner{ctx: ctx, inner: sp.Close}
 }
 
 func (s *spinner) ok(msg string)   { s.done("✓", msg) }
 func (s *spinner) fail(msg string) { s.done("✗", msg) }
 
+// Close stops the spinner without printing a marker. Safe to call multiple
+// times — combine with `defer s.Close()` to guarantee cleanup on early
+// returns.
+func (s *spinner) Close() { s.done("", "") }
+
 func (s *spinner) done(mark, msg string) {
 	if s.finished {
 		return
 	}
 	s.finished = true
-	s.close()
-	cmdio.LogString(s.ctx, "  "+cmdio.Cyan(s.ctx, mark)+" "+msg)
+	s.inner()
+	if mark != "" {
+		cmdio.LogString(s.ctx, "  "+cmdio.Cyan(s.ctx, mark)+" "+msg)
+	}
 }
 
 // status formats a lakebox lifecycle status with a color hint.

From d356344e1de8d0cbf785893dee958b60e2503196 Mon Sep 17 00:00:00 2001
From: Pieter Noordhuis <pieter.noordhuis@databricks.com>
Date: Thu, 7 May 2026 11:09:22 +0200
Subject: [PATCH 24/36] lakebox: hold cmdio spinner via interface; drop
 redundant 'finished' gate
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Define a local cmdioSpinner interface (just Close()) that the unexported
cmdio.spinner type satisfies structurally. Embed it on our wrapper so
spinner.Close comes for free, and drop the func() workaround.

The 'finished' bool was only preventing double-printing the ✓/✗ marker
if a caller called ok/fail twice — caller pilot error rather than a real
hazard, and cmdio's own Close is already idempotent (sync.OnceFunc on
sendQuit), so the gate isn't needed for resource safety. Net effect:
shorter, the embedded Close() is still safe to defer, and double-calls
to ok/fail print twice (which they always should have).

Co-authored-by: Isaac
---
 cmd/lakebox/ui.go | 45 +++++++++++++++++++--------------------------
 1 file changed, 19 insertions(+), 26 deletions(-)

diff --git a/cmd/lakebox/ui.go b/cmd/lakebox/ui.go
index c2669ce8ef..29a7d1274d 100644
--- a/cmd/lakebox/ui.go
+++ b/cmd/lakebox/ui.go
@@ -9,41 +9,34 @@ import (
 	"github.com/databricks/cli/libs/cmdio"
 )
 
-// spinner wraps cmdio.NewSpinner with terminal ok/fail markers. The first
-// call to ok, fail, or Close closes the underlying cmdio spinner; ok/fail
-// also log a final line to stderr. Subsequent calls are no-ops, so callers
-// are expected to `defer s.Close()` and call ok/fail on the success/failure
-// path. Close on its own (no marker) just stops the spinner — useful when an
-// error path returns before reaching ok/fail.
+// cmdioSpinner is the subset of *cmdio.spinner's method set we need.
+// Defining the interface locally lets us hold the unexported type as a
+// struct field; cmdio's spinner satisfies it structurally.
+type cmdioSpinner interface {
+	Close()
+}
+
+// spinner wraps cmdio.NewSpinner with ok/fail markers. ok and fail close the
+// underlying spinner and log a final ✓/✗ line; Close stops the spinner
+// without printing. cmdio's Close is itself idempotent, so a `defer s.Close()`
+// is safe alongside an ok/fail call on the success path.
 type spinner struct {
-	ctx      context.Context
-	inner    func()
-	finished bool
+	cmdioSpinner
+	ctx context.Context
 }
 
 func spin(ctx context.Context, msg string) *spinner {
 	sp := cmdio.NewSpinner(ctx)
 	sp.Update(msg)
-	return &spinner{ctx: ctx, inner: sp.Close}
+	return &spinner{cmdioSpinner: sp, ctx: ctx}
 }
 
-func (s *spinner) ok(msg string)   { s.done("✓", msg) }
-func (s *spinner) fail(msg string) { s.done("✗", msg) }
-
-// Close stops the spinner without printing a marker. Safe to call multiple
-// times — combine with `defer s.Close()` to guarantee cleanup on early
-// returns.
-func (s *spinner) Close() { s.done("", "") }
+func (s *spinner) ok(msg string)   { s.mark("✓", msg) }
+func (s *spinner) fail(msg string) { s.mark("✗", msg) }
 
-func (s *spinner) done(mark, msg string) {
-	if s.finished {
-		return
-	}
-	s.finished = true
-	s.inner()
-	if mark != "" {
-		cmdio.LogString(s.ctx, "  "+cmdio.Cyan(s.ctx, mark)+" "+msg)
-	}
+func (s *spinner) mark(mark, msg string) {
+	s.Close()
+	cmdio.LogString(s.ctx, "  "+cmdio.Cyan(s.ctx, mark)+" "+msg)
 }
 
 // status formats a lakebox lifecycle status with a color hint.

From e6e461f4ae939e8c3bd275bd6e21ac90d2b70da6 Mon Sep 17 00:00:00 2001
From: Pieter Noordhuis <pieter.noordhuis@databricks.com>
Date: Thu, 7 May 2026 11:10:10 +0200
Subject: [PATCH 25/36] lakebox: expose Update through the spinner wrapper

Add Update(msg string) to the cmdioSpinner interface so callers can
re-suffix the spinner mid-spin without reaching past our wrapper. No
current call site uses it, but it's a free pass-through via embedding
and matches the underlying cmdio API.

Co-authored-by: Isaac
---
 cmd/lakebox/ui.go | 1 +
 1 file changed, 1 insertion(+)

diff --git a/cmd/lakebox/ui.go b/cmd/lakebox/ui.go
index 29a7d1274d..a2904c7fe2 100644
--- a/cmd/lakebox/ui.go
+++ b/cmd/lakebox/ui.go
@@ -13,6 +13,7 @@ import (
 // Defining the interface locally lets us hold the unexported type as a
 // struct field; cmdio's spinner satisfies it structurally.
 type cmdioSpinner interface {
+	Update(msg string)
 	Close()
 }
 

From adb7d73b64f3059050d6777a6f39d0924d29e9dc Mon Sep 17 00:00:00 2001
From: Pieter Noordhuis <pieter.noordhuis@databricks.com>
Date: Thu, 7 May 2026 11:44:28 +0200
Subject: [PATCH 26/36] lakebox: validate saved default before reusing it on
 ssh
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

If the default lakebox stored at ~/.databricks/lakebox.json gets removed
out-of-band (auto-stop expiry, admin reap, deletion from another
machine), 'lakebox ssh' would happily try to ssh to it via the gateway
and the user would get a confusing 'Permission denied (publickey)' from
ssh. There was no signal that the default was stale.

api.get the saved default first; if it 404s (or any other error), warn,
clearDefault, and fall through to the existing 'no default → provision a
fresh one' branch. Mirrors the validation already in 'lakebox create'.

Co-authored-by: Isaac
---
 cmd/lakebox/ssh.go | 23 ++++++++++++++++++-----
 1 file changed, 18 insertions(+), 5 deletions(-)

diff --git a/cmd/lakebox/ssh.go b/cmd/lakebox/ssh.go
index 2bbd5a7e34..9ba1b4957a 100644
--- a/cmd/lakebox/ssh.go
+++ b/cmd/lakebox/ssh.go
@@ -88,13 +88,26 @@ Examples:
 
 			// Determine lakebox ID if not explicit.
 			if lakeboxID == "" {
+				api, err := newLakeboxAPI(w)
+				if err != nil {
+					return err
+				}
+
+				// If we have a saved default, confirm it still exists on the
+				// server. The lakebox may have been auto-stopped, deleted from
+				// another machine, or reaped by an admin since we wrote the
+				// state file. Clear the stale entry and fall through to
+				// provisioning a fresh one.
 				if def := getDefault(ctx, profile); def != "" {
-					lakeboxID = def
-				} else {
-					api, err := newLakeboxAPI(w)
-					if err != nil {
-						return err
+					if _, err := api.get(ctx, def); err == nil {
+						lakeboxID = def
+					} else {
+						warn(ctx, fmt.Sprintf("Saved default %s is gone; provisioning a new lakebox", def))
+						_ = clearDefault(ctx, profile)
 					}
+				}
+
+				if lakeboxID == "" {
 					pubKeyData, err := os.ReadFile(keyPath + ".pub")
 					if err != nil {
 						return fmt.Errorf("failed to read public key %s.pub: %w", keyPath, err)

From a6eece8941f4ff9ac431199dbfd39cca6d47a5e9 Mon Sep 17 00:00:00 2001
From: Pieter Noordhuis <pieter.noordhuis@databricks.com>
Date: Thu, 7 May 2026 11:44:29 +0200
Subject: [PATCH 27/36] lakebox: add unit tests for state.go
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Cover load/save/clear round-trips, missing-file and corrupt-JSON paths,
multi-profile independence, and the legacy 'last_profile' field that
older CLI versions wrote — loadState must accept it (silently dropping
the unknown key) and saveState must rewrite the file without it so it
naturally falls off on the next mutation.

All tests use env.WithUserHomeDir(t.Context(), t.TempDir()) so they
operate on an isolated state file.

Co-authored-by: Isaac
---
 cmd/lakebox/state_test.go | 138 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 138 insertions(+)
 create mode 100644 cmd/lakebox/state_test.go

diff --git a/cmd/lakebox/state_test.go b/cmd/lakebox/state_test.go
new file mode 100644
index 0000000000..9755117b3b
--- /dev/null
+++ b/cmd/lakebox/state_test.go
@@ -0,0 +1,138 @@
+package lakebox
+
+import (
+	"context"
+	"io/fs"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/databricks/cli/libs/env"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// stateCtx returns a context whose $HOME is a temp directory, so state file
+// operations are isolated from the developer's real ~/.databricks/lakebox.json.
+func stateCtx(t *testing.T) (context.Context, string) {
+	t.Helper()
+	home := t.TempDir()
+	ctx := env.WithUserHomeDir(t.Context(), home)
+	return ctx, filepath.Join(home, ".databricks", "lakebox.json")
+}
+
+func TestStateLoadMissingFileReturnsEmpty(t *testing.T) {
+	ctx, _ := stateCtx(t)
+	state, err := loadState(ctx)
+	require.NoError(t, err)
+	assert.Equal(t, &stateFile{Defaults: map[string]string{}}, state)
+}
+
+func TestStateGetDefaultMissingProfileReturnsEmpty(t *testing.T) {
+	ctx, _ := stateCtx(t)
+	assert.Equal(t, "", getDefault(ctx, "any-profile"))
+}
+
+func TestStateSetGetDefaultRoundTrip(t *testing.T) {
+	ctx, _ := stateCtx(t)
+
+	require.NoError(t, setDefault(ctx, "profile-a", "lakebox-a"))
+	assert.Equal(t, "lakebox-a", getDefault(ctx, "profile-a"))
+	assert.Equal(t, "", getDefault(ctx, "profile-b"))
+}
+
+func TestStateMultipleProfilesIndependent(t *testing.T) {
+	ctx, _ := stateCtx(t)
+
+	require.NoError(t, setDefault(ctx, "profile-a", "lakebox-a"))
+	require.NoError(t, setDefault(ctx, "profile-b", "lakebox-b"))
+
+	assert.Equal(t, "lakebox-a", getDefault(ctx, "profile-a"))
+	assert.Equal(t, "lakebox-b", getDefault(ctx, "profile-b"))
+}
+
+func TestStateSetDefaultOverwrites(t *testing.T) {
+	ctx, _ := stateCtx(t)
+
+	require.NoError(t, setDefault(ctx, "profile-a", "lakebox-a"))
+	require.NoError(t, setDefault(ctx, "profile-a", "lakebox-a-prime"))
+	assert.Equal(t, "lakebox-a-prime", getDefault(ctx, "profile-a"))
+}
+
+func TestStateClearDefault(t *testing.T) {
+	ctx, _ := stateCtx(t)
+
+	require.NoError(t, setDefault(ctx, "profile-a", "lakebox-a"))
+	require.NoError(t, setDefault(ctx, "profile-b", "lakebox-b"))
+
+	require.NoError(t, clearDefault(ctx, "profile-a"))
+	assert.Equal(t, "", getDefault(ctx, "profile-a"))
+	assert.Equal(t, "lakebox-b", getDefault(ctx, "profile-b"))
+}
+
+func TestStateClearDefaultMissingProfileNoError(t *testing.T) {
+	ctx, _ := stateCtx(t)
+	assert.NoError(t, clearDefault(ctx, "no-such-profile"))
+}
+
+// Pre-existing files from earlier CLI versions carry a `last_profile` field
+// the current schema doesn't know about. loadState must accept the file
+// (silently dropping the unknown field) and saveState must rewrite without
+// it, so the field naturally falls off on the next mutation.
+func TestStateLoadIgnoresUnknownFields(t *testing.T) {
+	ctx, path := stateCtx(t)
+	require.NoError(t, os.MkdirAll(filepath.Dir(path), 0o700))
+	require.NoError(t, os.WriteFile(path, []byte(`{
+        "defaults": {"profile-a": "lakebox-a"},
+        "last_profile": "profile-a"
+    }`), 0o600))
+
+	assert.Equal(t, "lakebox-a", getDefault(ctx, "profile-a"))
+
+	require.NoError(t, setDefault(ctx, "profile-a", "lakebox-a-prime"))
+	rewritten, err := os.ReadFile(path)
+	require.NoError(t, err)
+	assert.NotContains(t, string(rewritten), "last_profile")
+}
+
+func TestStateLoadReturnsErrorOnCorruptJSON(t *testing.T) {
+	ctx, path := stateCtx(t)
+	require.NoError(t, os.MkdirAll(filepath.Dir(path), 0o700))
+	require.NoError(t, os.WriteFile(path, []byte("{not valid json"), 0o600))
+
+	_, err := loadState(ctx)
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "failed to parse")
+}
+
+// Files written by saveState must round-trip through loadState even if the
+// caller starts from an empty Defaults map.
+func TestStateSaveCreatesParentDirs(t *testing.T) {
+	ctx, path := stateCtx(t)
+
+	// Confirm parent dir does not exist yet.
+	_, err := os.Stat(filepath.Dir(path))
+	assert.ErrorIs(t, err, fs.ErrNotExist)
+
+	require.NoError(t, setDefault(ctx, "profile-a", "lakebox-a"))
+
+	// File and parent dir now exist with sensible perms.
+	info, err := os.Stat(path)
+	require.NoError(t, err)
+	assert.Equal(t, os.FileMode(0o600), info.Mode().Perm())
+
+	dirInfo, err := os.Stat(filepath.Dir(path))
+	require.NoError(t, err)
+	assert.Equal(t, os.FileMode(0o700), dirInfo.Mode().Perm())
+}
+
+// Defaults of nil on disk (legal but not what saveState produces) must still
+// load to a usable empty map so callers can setDefault without nil-deref.
+func TestStateLoadNilDefaultsMap(t *testing.T) {
+	ctx, path := stateCtx(t)
+	require.NoError(t, os.MkdirAll(filepath.Dir(path), 0o700))
+	require.NoError(t, os.WriteFile(path, []byte(`{}`), 0o600))
+
+	require.NoError(t, setDefault(ctx, "profile-a", "lakebox-a"))
+	assert.Equal(t, "lakebox-a", getDefault(ctx, "profile-a"))
+}

From bd72f850ce5774ec174719119499da8289dab5cf Mon Sep 17 00:00:00 2001
From: Pieter Noordhuis <pieter.noordhuis@databricks.com>
Date: Thu, 7 May 2026 14:41:54 +0200
Subject: [PATCH 28/36] lakebox: add keyHash helper matching the server's
 algorithm
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The /api/2.0/lakebox/ssh-keys endpoint identifies registered keys by
hash. Live exploration confirmed the algorithm: sha256 of
'<type> <base64-blob>' (comment stripped) truncated to 16 bytes, hex
encoded — looks like MD5 (32 hex chars) but isn't.

Encode this client-side so we can answer 'is this local key registered?'
without a list call. Tests use the exact hashes captured from the live
API as ground truth, plus an edge case for empty input.

Co-authored-by: Isaac
---
 cmd/lakebox/keyhash.go      | 27 +++++++++++++++++
 cmd/lakebox/keyhash_test.go | 58 +++++++++++++++++++++++++++++++++++++
 2 files changed, 85 insertions(+)
 create mode 100644 cmd/lakebox/keyhash.go
 create mode 100644 cmd/lakebox/keyhash_test.go

diff --git a/cmd/lakebox/keyhash.go b/cmd/lakebox/keyhash.go
new file mode 100644
index 0000000000..7bb9fc3191
--- /dev/null
+++ b/cmd/lakebox/keyhash.go
@@ -0,0 +1,27 @@
+package lakebox
+
+import (
+	"crypto/sha256"
+	"encoding/hex"
+	"strings"
+)
+
+// keyHash returns the identifier the lakebox SSH-keys API assigns to a
+// public key. The algorithm is sha256("<type> <base64-blob>") truncated to
+// the first 16 bytes and hex-encoded; the OpenSSH comment (anything after
+// the second whitespace-separated token) is stripped before hashing, so
+// registering the same key under different comments yields the same hash.
+// Inputs that don't have a second token are hashed as-is.
+//
+// Useful for client-side checks like "is the local lakebox_rsa.pub already
+// registered?" without a list call against /api/2.0/lakebox/ssh-keys.
+func keyHash(publicKey string) string {
+	canonical := publicKey
+	if i := strings.IndexByte(publicKey, ' '); i >= 0 {
+		if j := strings.IndexByte(publicKey[i+1:], ' '); j >= 0 {
+			canonical = publicKey[:i+1+j]
+		}
+	}
+	sum := sha256.Sum256([]byte(canonical))
+	return hex.EncodeToString(sum[:16])
+}
diff --git a/cmd/lakebox/keyhash_test.go b/cmd/lakebox/keyhash_test.go
new file mode 100644
index 0000000000..4f40a42c78
--- /dev/null
+++ b/cmd/lakebox/keyhash_test.go
@@ -0,0 +1,58 @@
+package lakebox
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+// All expected hashes were captured live from /api/2.0/lakebox/ssh-keys
+// (see PR description); they're the ground truth for the algorithm.
+func TestKeyHash(t *testing.T) {
+	tests := []struct {
+		name  string
+		input string
+		want  string
+	}{
+		{
+			name:  "single-token input hashed verbatim",
+			input: "a",
+			want:  "ca978112ca1bbdcafac231b39a23dc4d",
+		},
+		{
+			name:  "type and blob with no comment",
+			input: "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIDUMMY",
+			want:  "2b366430eb9743668b652921d3b22d54",
+		},
+		{
+			name:  "comment is stripped before hashing",
+			input: "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIDUMMY comment-one",
+			want:  "2b366430eb9743668b652921d3b22d54",
+		},
+		{
+			name:  "different comment same key still matches",
+			input: "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIDUMMY entirely-different-comment",
+			want:  "2b366430eb9743668b652921d3b22d54",
+		},
+		{
+			name:  "longer key with multi-word comment",
+			input: "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAITESTKEY1234 test-from-cli-exploration",
+			want:  "52c927705154e2d98a1b7036cc3e06dc",
+		},
+		{
+			name:  "empty input still produces a hash",
+			input: "",
+			want:  "e3b0c44298fc1c149afbf4c8996fb924",
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			assert.Equal(t, tc.want, keyHash(tc.input))
+		})
+	}
+}
+
+func TestKeyHashIsStableLength(t *testing.T) {
+	// 16 bytes hex-encoded = 32 chars, matching what the API returns.
+	assert.Len(t, keyHash("anything"), 32)
+}

From 4d4ca9e37ff04ac0a84f324c4ef7568e17ec0bd8 Mon Sep 17 00:00:00 2001
From: Pieter Noordhuis <pieter.noordhuis@databricks.com>
Date: Thu, 7 May 2026 14:45:45 +0200
Subject: [PATCH 29/36] lakebox: simplify keyHash with strings.SplitSeq;
 correct doc comment

Two small fixes to the keyHash helper:

- Replace the nested IndexByte calls with a range over strings.SplitSeq
  that breaks after the second token. Tracks a running byte offset so we
  still slice the original string instead of allocating a joined copy.
- Drop the misleading 'without a list call' phrasing. You still need to
  call GET /ssh-keys; the helper just means you can match a local key
  against the listing by hash, without re-uploading the key contents.

Co-authored-by: Isaac
---
 cmd/lakebox/keyhash.go | 23 ++++++++++++++++-------
 1 file changed, 16 insertions(+), 7 deletions(-)

diff --git a/cmd/lakebox/keyhash.go b/cmd/lakebox/keyhash.go
index 7bb9fc3191..f401e255a7 100644
--- a/cmd/lakebox/keyhash.go
+++ b/cmd/lakebox/keyhash.go
@@ -13,15 +13,24 @@ import (
 // registering the same key under different comments yields the same hash.
 // Inputs that don't have a second token are hashed as-is.
 //
-// Useful for client-side checks like "is the local lakebox_rsa.pub already
-// registered?" without a list call against /api/2.0/lakebox/ssh-keys.
+// Useful for matching a locally-known key against entries in a
+// GET /ssh-keys listing without sending the key contents back to the
+// server.
 func keyHash(publicKey string) string {
-	canonical := publicKey
-	if i := strings.IndexByte(publicKey, ' '); i >= 0 {
-		if j := strings.IndexByte(publicKey[i+1:], ' '); j >= 0 {
-			canonical = publicKey[:i+1+j]
+	// Walk the splits and break out after the second token; the
+	// running offset is what we slice the original string by.
+	end := 0
+	seen := 0
+	for token := range strings.SplitSeq(publicKey, " ") {
+		if seen > 0 {
+			end++ // separator before this token
+		}
+		end += len(token)
+		seen++
+		if seen == 2 {
+			break
 		}
 	}
-	sum := sha256.Sum256([]byte(canonical))
+	sum := sha256.Sum256([]byte(publicKey[:end]))
 	return hex.EncodeToString(sum[:16])
 }

From 9b696ba6f9806cb4b00b46d10408fccbada88523 Mon Sep 17 00:00:00 2001
From: Pieter Noordhuis <pieter.noordhuis@databricks.com>
Date: Thu, 7 May 2026 14:47:44 +0200
Subject: [PATCH 30/36] lakebox: simplify keyHash to byte iteration

The SplitSeq approach needed a running offset and a 'first iteration?'
guard inside the loop. Walking bytes directly until we see the second
space is shorter and reads more directly: count spaces, slice when the
counter hits 2. Single pass, no allocation.

Co-authored-by: Isaac
---
 cmd/lakebox/keyhash.go | 23 ++++++++++-------------
 1 file changed, 10 insertions(+), 13 deletions(-)

diff --git a/cmd/lakebox/keyhash.go b/cmd/lakebox/keyhash.go
index f401e255a7..7f4fcd0bd4 100644
--- a/cmd/lakebox/keyhash.go
+++ b/cmd/lakebox/keyhash.go
@@ -3,7 +3,6 @@ package lakebox
 import (
 	"crypto/sha256"
 	"encoding/hex"
-	"strings"
 )
 
 // keyHash returns the identifier the lakebox SSH-keys API assigns to a
@@ -17,18 +16,16 @@ import (
 // GET /ssh-keys listing without sending the key contents back to the
 // server.
 func keyHash(publicKey string) string {
-	// Walk the splits and break out after the second token; the
-	// running offset is what we slice the original string by.
-	end := 0
-	seen := 0
-	for token := range strings.SplitSeq(publicKey, " ") {
-		if seen > 0 {
-			end++ // separator before this token
-		}
-		end += len(token)
-		seen++
-		if seen == 2 {
-			break
+	// Slice off the OpenSSH comment by stopping at the second space.
+	end := len(publicKey)
+	spaces := 0
+	for i, c := range publicKey {
+		if c == ' ' {
+			spaces++
+			if spaces == 2 {
+				end = i
+				break
+			}
 		}
 	}
 	sum := sha256.Sum256([]byte(publicKey[:end]))

From 34aaad65eed455a5f58500cdc9b9afbf8384dff3 Mon Sep 17 00:00:00 2001
From: Pieter Noordhuis <pieter.noordhuis@databricks.com>
Date: Thu, 7 May 2026 14:49:20 +0200
Subject: [PATCH 31/36] lakebox: correct misleading comment on keyHash test
 inputs
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The previous wording implied the expected hashes were pulled from real
registered keys returned by the API. They aren't — they're sha256[:16]
of synthetic strings I posted during exploration. The algorithm was
verified live; the test pins the algorithm rather than any specific
captured registration.

Co-authored-by: Isaac
---
 cmd/lakebox/keyhash_test.go | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/cmd/lakebox/keyhash_test.go b/cmd/lakebox/keyhash_test.go
index 4f40a42c78..a03bb1a209 100644
--- a/cmd/lakebox/keyhash_test.go
+++ b/cmd/lakebox/keyhash_test.go
@@ -6,8 +6,10 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-// All expected hashes were captured live from /api/2.0/lakebox/ssh-keys
-// (see PR description); they're the ground truth for the algorithm.
+// Inputs are synthetic; expected values are sha256(canonical input)[:16]
+// in hex. The algorithm was verified against the live
+// /api/2.0/lakebox/ssh-keys endpoint during exploration, so this test
+// pins the algorithm — not a known set of real registered keys.
 func TestKeyHash(t *testing.T) {
 	tests := []struct {
 		name  string

From 4df1daf4eead62236944f9a44a22856ef5dbb06b Mon Sep 17 00:00:00 2001
From: Pieter Noordhuis <pieter.noordhuis@databricks.com>
Date: Thu, 7 May 2026 14:53:29 +0200
Subject: [PATCH 32/36] lakebox: drop superfluous TestKeyHashIsStableLength

Every case in TestKeyHash already pins an exact 32-char hex string, so a
separate length-only test buys nothing.

Co-authored-by: Isaac
---
 cmd/lakebox/keyhash_test.go | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/cmd/lakebox/keyhash_test.go b/cmd/lakebox/keyhash_test.go
index a03bb1a209..638f1d8f34 100644
--- a/cmd/lakebox/keyhash_test.go
+++ b/cmd/lakebox/keyhash_test.go
@@ -53,8 +53,3 @@ func TestKeyHash(t *testing.T) {
 		})
 	}
 }
-
-func TestKeyHashIsStableLength(t *testing.T) {
-	// 16 bytes hex-encoded = 32 chars, matching what the API returns.
-	assert.Len(t, keyHash("anything"), 32)
-}

From 670f66e749a16df5114894ea4744f0ea38d6df81 Mon Sep 17 00:00:00 2001
From: Pieter Noordhuis <pieter.noordhuis@databricks.com>
Date: Fri, 8 May 2026 15:46:59 +0200
Subject: [PATCH 33/36] lakebox: align org-ID header with the rest of the
 codebase

Drop the bespoke resolveWorkspaceID helper and the cached wsID field on
lakeboxAPI. Match the minimal pattern that libs/telemetry, libs/filer,
and SDK-generated workspace services already use: read cfg.WorkspaceID
directly, send the X-Databricks-Org-Id header if set.

Removes the '?o=<id>' fallback that parsed the host's query string.
That behavior was unique to lakebox and inconsistent with how every
other CLI surface handles SPOG hosts; the SDK's host-metadata discovery
populates cfg.WorkspaceID for hosts that need it, and users who run
into edge cases set workspace_id explicitly the same way they would
for `bundle deploy` or `databricks api`.

Adds the auth.WorkspaceIDNone ("none") sentinel strip so a profile
created via `databricks auth login` for SPOG account-level access
doesn't send the literal string "none" as the routing identifier.
This fix matches what cmd/api/api.go (#5137) and libs/auth do; the
four other orgIDHeaders helpers in the codebase still have the latent
bug, which is a separate cleanup.

Co-authored-by: Isaac
---
 cmd/lakebox/api.go | 32 +++++++++++---------------------
 1 file changed, 11 insertions(+), 21 deletions(-)

diff --git a/cmd/lakebox/api.go b/cmd/lakebox/api.go
index 2b65fd8c70..754da218ec 100644
--- a/cmd/lakebox/api.go
+++ b/cmd/lakebox/api.go
@@ -4,10 +4,10 @@ import (
 	"context"
 	"fmt"
 	"net/http"
-	"net/url"
 	"strings"
 	"time"
 
+	"github.com/databricks/cli/libs/auth"
 	"github.com/databricks/databricks-sdk-go"
 	"github.com/databricks/databricks-sdk-go/client"
 )
@@ -28,8 +28,7 @@ const orgIDHeader = "X-Databricks-Org-Id"
 
 // lakeboxAPI wraps the SDK ApiClient with workspace-id-aware request headers.
 type lakeboxAPI struct {
-	c    *client.DatabricksClient
-	wsID string
+	c *client.DatabricksClient
 }
 
 // createRequest is the JSON body for POST /api/2.0/lakebox/sandboxes.
@@ -163,29 +162,20 @@ func newLakeboxAPI(w *databricks.WorkspaceClient) (*lakeboxAPI, error) {
 	if err != nil {
 		return nil, fmt.Errorf("failed to create lakebox API client: %w", err)
 	}
-	return &lakeboxAPI{c: c, wsID: resolveWorkspaceID(w)}, nil
-}
-
-// resolveWorkspaceID returns the workspace ID for the org-id header. Falls
-// back to the `?o=<id>` query parameter on the Host because some staging
-// gateways are configured that way and the SDK does not strip it into
-// Config.WorkspaceID.
-func resolveWorkspaceID(w *databricks.WorkspaceClient) string {
-	if id := w.Config.WorkspaceID; id != "" {
-		return id
-	}
-	parsed, err := url.Parse(w.Config.Host)
-	if err != nil {
-		return ""
-	}
-	return parsed.Query().Get("o")
+	return &lakeboxAPI{c: c}, nil
 }
 
+// headers attaches the workspace routing identifier so multi-workspace
+// gateways (e.g. SPOG hosts) can scope the credential. Mirrors the pattern
+// in libs/telemetry, libs/filer, and SDK-generated workspace services. The
+// auth.WorkspaceIDNone sentinel ("none") is treated as unset so the literal
+// string never goes on the wire.
 func (a *lakeboxAPI) headers() map[string]string {
-	if a.wsID == "" {
+	wsID := a.c.Config.WorkspaceID
+	if wsID == "" || wsID == auth.WorkspaceIDNone {
 		return nil
 	}
-	return map[string]string{orgIDHeader: a.wsID}
+	return map[string]string{orgIDHeader: wsID}
 }
 
 // create calls POST /api/2.0/lakebox/sandboxes with an optional public key.

From f6f28ebcb586688d22220f4dab324d89ea1a21ff Mon Sep 17 00:00:00 2001
From: Pieter Noordhuis <pieter.noordhuis@databricks.com>
Date: Fri, 8 May 2026 16:26:54 +0200
Subject: [PATCH 34/36] lakebox: skip state file writes when nothing changed

setDefault and clearDefault unconditionally rewrote ~/.databricks/
lakebox.json even when the in-memory state was identical to what was
already on disk: clearing a profile that wasn't in the map, or
re-setting the same value. That created or touched the file for
no-op operations.

Add change-detection guards to both: setDefault is a no-op when the
profile already maps to the requested ID; clearDefault is a no-op
when the profile isn't in the map. Result: a CLI invocation that
doesn't change state can no longer cause a file to spring into
existence on a fresh machine.

Tests:
 - clearDefault on a missing profile must leave the file absent
 - setDefault with an unchanged value must not bump the mtime
 - getDefault on a fresh state must not create the file (regression
   test for the read-only path)

Co-authored-by: Isaac
---
 cmd/lakebox/state.go      |  6 ++++++
 cmd/lakebox/state_test.go | 33 ++++++++++++++++++++++++++++++---
 2 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/cmd/lakebox/state.go b/cmd/lakebox/state.go
index 87cc96e78f..5be3da1d4a 100644
--- a/cmd/lakebox/state.go
+++ b/cmd/lakebox/state.go
@@ -81,6 +81,9 @@ func setDefault(ctx context.Context, profile, lakeboxID string) error {
 	if err != nil {
 		return err
 	}
+	if state.Defaults[profile] == lakeboxID {
+		return nil
+	}
 	state.Defaults[profile] = lakeboxID
 	return saveState(ctx, state)
 }
@@ -90,6 +93,9 @@ func clearDefault(ctx context.Context, profile string) error {
 	if err != nil {
 		return err
 	}
+	if _, ok := state.Defaults[profile]; !ok {
+		return nil
+	}
 	delete(state.Defaults, profile)
 	return saveState(ctx, state)
 }
diff --git a/cmd/lakebox/state_test.go b/cmd/lakebox/state_test.go
index 9755117b3b..2f7f591392 100644
--- a/cmd/lakebox/state_test.go
+++ b/cmd/lakebox/state_test.go
@@ -70,9 +70,36 @@ func TestStateClearDefault(t *testing.T) {
 	assert.Equal(t, "lakebox-b", getDefault(ctx, "profile-b"))
 }
 
-func TestStateClearDefaultMissingProfileNoError(t *testing.T) {
-	ctx, _ := stateCtx(t)
-	assert.NoError(t, clearDefault(ctx, "no-such-profile"))
+func TestStateClearDefaultMissingProfileDoesNotCreateFile(t *testing.T) {
+	ctx, path := stateCtx(t)
+
+	require.NoError(t, clearDefault(ctx, "no-such-profile"))
+
+	_, err := os.Stat(path)
+	assert.ErrorIs(t, err, fs.ErrNotExist, "clearDefault must not create the state file when there's nothing to remove")
+}
+
+func TestStateSetDefaultSameValueDoesNotRewriteFile(t *testing.T) {
+	ctx, path := stateCtx(t)
+
+	require.NoError(t, setDefault(ctx, "profile-a", "lakebox-a"))
+	before, err := os.Stat(path)
+	require.NoError(t, err)
+
+	// Re-set with the same value should be a no-op.
+	require.NoError(t, setDefault(ctx, "profile-a", "lakebox-a"))
+	after, err := os.Stat(path)
+	require.NoError(t, err)
+	assert.Equal(t, before.ModTime(), after.ModTime(), "no-op setDefault must not rewrite the file")
+}
+
+func TestStateSetDefaultMissingNoFileBeforeWrite(t *testing.T) {
+	ctx, path := stateCtx(t)
+
+	// Loading state on a fresh tempdir must not create the file.
+	assert.Equal(t, "", getDefault(ctx, "profile-a"))
+	_, err := os.Stat(path)
+	assert.ErrorIs(t, err, fs.ErrNotExist, "getDefault must not create the state file")
 }
 
 // Pre-existing files from earlier CLI versions carry a `last_profile` field

From 5807f245e2e07216122f26d36e155b80fb45885c Mon Sep 17 00:00:00 2001
From: Pieter Noordhuis <pieter.noordhuis@databricks.com>
Date: Fri, 8 May 2026 16:52:44 +0200
Subject: [PATCH 35/36] lakebox: hide the subcommand from the top-level help
 listing
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Mark the lakebox parent command as Hidden so it doesn't appear in
'databricks --help' under Developer Tools. The subcommands themselves
are still reachable — 'databricks lakebox --help' lists them — but the
feature stays out of the discoverable surface while it remains
internal.

This also reverts the acceptance/help/output.txt regen from the
previous push, since hiding the command means the golden file already
matches the actual help output.

Co-authored-by: Isaac
---
 cmd/lakebox/lakebox.go | 1 +
 1 file changed, 1 insertion(+)

diff --git a/cmd/lakebox/lakebox.go b/cmd/lakebox/lakebox.go
index b6c2970760..c4f7b6cc7e 100644
--- a/cmd/lakebox/lakebox.go
+++ b/cmd/lakebox/lakebox.go
@@ -9,6 +9,7 @@ func New() *cobra.Command {
 		Use:     "lakebox",
 		Short:   "Manage Databricks Lakebox environments",
 		GroupID: "development",
+		Hidden:  true,
 		Long: `Manage Databricks Lakebox environments.
 
 Lakebox provides SSH-accessible development environments backed by

From 33f949d2979adb8667e146e211b9effba03ddee9 Mon Sep 17 00:00:00 2001
From: Pieter Noordhuis <pieter.noordhuis@databricks.com>
Date: Mon, 11 May 2026 17:35:23 +0200
Subject: [PATCH 36/36] lakebox: skip Unix perm assertions in state test on
 Windows

Go on Windows synthesizes file mode from the read-only attribute
(0o666/0o777), so the 0o600/0o700 assertions can never hold. Matches the
existing pattern used in libs/cache, libs/completion, and the ssh vscode
settings tests.

Co-authored-by: Isaac
---
 cmd/lakebox/state_test.go | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/cmd/lakebox/state_test.go b/cmd/lakebox/state_test.go
index 2f7f591392..a7488de2e1 100644
--- a/cmd/lakebox/state_test.go
+++ b/cmd/lakebox/state_test.go
@@ -5,6 +5,7 @@ import (
 	"io/fs"
 	"os"
 	"path/filepath"
+	"runtime"
 	"testing"
 
 	"github.com/databricks/cli/libs/env"
@@ -146,11 +147,16 @@ func TestStateSaveCreatesParentDirs(t *testing.T) {
 	// File and parent dir now exist with sensible perms.
 	info, err := os.Stat(path)
 	require.NoError(t, err)
-	assert.Equal(t, os.FileMode(0o600), info.Mode().Perm())
 
 	dirInfo, err := os.Stat(filepath.Dir(path))
 	require.NoError(t, err)
-	assert.Equal(t, os.FileMode(0o700), dirInfo.Mode().Perm())
+
+	// Windows does not honor Unix permission bits; os.Stat reports 0o666/0o777
+	// regardless of what was passed to OpenFile/MkdirAll.
+	if runtime.GOOS != "windows" {
+		assert.Equal(t, os.FileMode(0o600), info.Mode().Perm())
+		assert.Equal(t, os.FileMode(0o700), dirInfo.Mode().Perm())
+	}
 }
 
 // Defaults of nil on disk (legal but not what saveState produces) must still