[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

Author: shuochen0311

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")
 	}

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