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/16] 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 014471f7638..fe81149c083 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 00000000000..ff8f7d30b12
--- /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 00000000000..872776cc8d5
--- /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 00000000000..9d5a366c9cd
--- /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 00000000000..a814083ed39
--- /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 00000000000..d47f629572b
--- /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 00000000000..6523debef91
--- /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 00000000000..bf80a9919e5
--- /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 00000000000..1978dec684e
--- /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 00000000000..c0c8ad2d84d
--- /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 00000000000..1afd968211d
--- /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/16] 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 ff8f7d30b12..94877b4a424 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 6523debef91..aa9463bca8c 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 bf80a9919e5..90139d6be8b 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 00000000000..5a19cc4f57a
--- /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 1afd968211d..4bb130496db 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/16] 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 fe81149c083..c120f25aa71 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 aa9463bca8c..127b5d93bfc 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 00000000000..7286a14bf5e
--- /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 5a19cc4f57a..00000000000
--- 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 1978dec684e..8868f38e81b 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/16] 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 c120f25aa71..ddbb70f4519 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 127b5d93bfc..4afa321241c 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 7286a14bf5e..a1da60422ba 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/16] 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 8868f38e81b..7559893bfbc 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/16] 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 ddbb70f4519..8a703755145 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 872776cc8d5..db1a22ebb7b 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 9d5a366c9cd..b632c5984af 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 a814083ed39..9c8ce939634 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 4afa321241c..6a968df87ac 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 90139d6be8b..3222d1c10c5 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 a1da60422ba..27d6cc59a16 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 8868f38e81b..86098baf5ab 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 c0c8ad2d84d..b84b5b16e1f 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 4bb130496db..eaeeb8d7ccf 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 00000000000..acdf259b4c9
--- /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 00000000000..c55c0aa182d
--- /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/16] 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 db1a22ebb7b..c4ce3a439ea 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 9c8ce939634..ba56e2a508d 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 3222d1c10c5..2ed3149658e 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 27d6cc59a16..f3550d8e5de 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 04a999bd404..483dbd38a8e 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 eaeeb8d7ccf..bf2efbcaba1 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 00000000000..2eab33310c4
--- /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/16] 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 94877b4a424..04cbc1179c6 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 2ed3149658e..69f9b2e3d7c 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 bf2efbcaba1..d362143dc67 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/16] 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 04cbc1179c6..06b6de217bd 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 c4ce3a439ea..096df26ce6b 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 69f9b2e3d7c..fe303028a00 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 483dbd38a8e..11297f27868 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 d362143dc67..aa4a443d0a0 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/16] 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 06b6de217bd..3d0a4707b42 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 fe303028a00..f058524e7ee 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 aa4a443d0a0..f5df1ee4a40 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/16] 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 3d0a4707b42..73d4b5d81f1 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 00000000000..e909489f2d2
--- /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 6a968df87ac..25a9b479e5b 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/16] =?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 73d4b5d81f1..288e704a9ef 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 e909489f2d2..fe3b80ddf29 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/16] 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 288e704a9ef..4fb3af26304 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/16] [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 4fb3af26304..acaeff47e8b 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 11297f27868..2a7db87a1b7 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 26672b60efb7810068c6b4730a7333e8170f1b5b Mon Sep 17 00:00:00 2001
From: shuochen0311 <shuo.chen@databricks.com>
Date: Wed, 20 May 2026 17:00:43 +0000
Subject: [PATCH 15/16] [lakebox] Default staging SSH gateway to ue1.s.dbrx.dev

us-east-1 is the staging region where the SSH gateway is reachable in
practice; uw2.s.dbrx.dev does not have a routable lakebox listener.
Users targeting a different staging region can still override with
--gateway.

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

diff --git a/cmd/lakebox/ssh.go b/cmd/lakebox/ssh.go
index 2a7db87a1b7..8d7bd5d51ea 100644
--- a/cmd/lakebox/ssh.go
+++ b/cmd/lakebox/ssh.go
@@ -13,7 +13,7 @@ import (
 
 const (
 	defaultGatewayHost        = "uw2.dbrx.dev"
-	stagingDefaultGatewayHost = "uw2.s.dbrx.dev"
+	stagingDefaultGatewayHost = "ue1.s.dbrx.dev"
 	defaultGatewayPort        = "2222"
 )
 

From e7d6235dc084938d6d87830358e967d83f6cd459 Mon Sep 17 00:00:00 2001
From: shuochen0311 <shuo.chen@databricks.com>
Date: Wed, 20 May 2026 17:52:16 +0000
Subject: [PATCH 16/16] [lakebox] Wrap create request body in {"sandbox": ...}

The gateway's CreateSandbox handler now rejects the legacy unwrapped
shape with INVALID_PARAMETER_VALUE: "unknown field \`public_key\`,
expected \`sandbox\`". Send the wrapped `{"sandbox": {}}` payload that
matches the `CreateSandboxRequest` proto. SSH keys are registered
separately via /api/2.0/lakebox/ssh-keys, so the body has nothing else
to carry today.

Verified end-to-end against staging us-east-1: create / list / ssh
through ue1.s.dbrx.dev all succeed.

Co-authored-by: Isaac
---
 cmd/lakebox/api.go | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/cmd/lakebox/api.go b/cmd/lakebox/api.go
index acaeff47e8b..6081e5abcc4 100644
--- a/cmd/lakebox/api.go
+++ b/cmd/lakebox/api.go
@@ -26,12 +26,11 @@ type lakeboxAPI struct {
 // 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.
+// field. Today every member is server-chosen so the payload is just an
+// empty wrapper; we still send the wrapper because the gateway no longer
+// accepts the unwrapped legacy shape.
 type createRequest struct {
-	PublicKey string `json:"public_key,omitempty"`
+	Sandbox struct{} `json:"sandbox"`
 }
 
 // createResponse is the JSON body returned by POST /api/2.0/lakebox/sandboxes.
@@ -140,9 +139,11 @@ 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}
+// create calls POST /api/2.0/lakebox/sandboxes. SSH keys are registered
+// separately via /api/2.0/lakebox/ssh-keys, so this body is just the empty
+// `Sandbox` wrapper.
+func (a *lakeboxAPI) create(ctx context.Context, _ string) (*createResponse, error) {
+	body := createRequest{}
 	jsonBody, err := json.Marshal(body)
 	if err != nil {
 		return nil, fmt.Errorf("failed to marshal request: %w", err)