Prompt for the host instead of failing when device code login has none

acceptance/cmd/auth/login/device-code/unsupported-host/output.txt

@@ -2,7 +2,7 @@
 === No host
 
 >>> musterr [CLI] auth login --profile new-profile --use-device-code
-Error: --use-device-code requires a workspace host. Pass it as an argument or with --host
+Error: the command is being run in a non-interactive environment, please specify a host using --host
 
 === Account host
 

acceptance/cmd/auth/login/device-code/unsupported-host/script

@@ -3,7 +3,8 @@ sethome "./home"
 title "No host\n"
 # The profile does not exist yet and no host is given. auth login does not read
 # DATABRICKS_HOST, so the usual login would fall back to the browser-based
-# discovery flow, which device code login does not support.
+# discovery flow. Device code login skips it and asks for the host instead,
+# which fails without a terminal.
 trace musterr $CLI auth login --profile new-profile --use-device-code
 
 envsubst < "./home/.databrickscfg.tmpl" > "./home/.databrickscfg"

cmd/auth/login.go

@@ -225,9 +225,15 @@ a new profile is created.
 		// the user re-login to an existing profile, create a new one, or enter
 		// a host URL. With no profiles configured the picker still shows the
 		// two action entries so the user can choose between web-based discovery
-		// (Create a new profile) and a manual host URL.
+		// (Create a new profile) and a manual host URL. Device code login skips
+		// discovery, so its "Create a new profile" asks for the host after the
+		// profile name, unless a profile with that name already has a host.
 		if profileName == "" && authArguments.Host == "" && len(args) == 0 && cmdio.IsPromptSupported(ctx) {
-			allProfiles, err := profile.DefaultProfiler.LoadProfiles(ctx, profile.MatchAllProfiles)
+			matchProfiles := profile.MatchAllProfiles
+			if useDeviceCode {
+				matchProfiles = matchDeviceCodeProfiles
+			}
+			allProfiles, err := profile.DefaultProfiler.LoadProfiles(ctx, matchProfiles)
 			if err != nil && !errors.Is(err, profile.ErrNoConfiguration) {
 				return err
 			}
@@ -252,6 +258,11 @@ a new profile is created.
 				if err != nil {
 					return err
 				}
+				// Fail here rather than prompting for the host again after
+				// the profile name.
+				if useDeviceCode && host == "" {
+					return errDeviceCodeNeedsHost
+				}
 				authArguments.Host = host
 			case profilePickerCreateNew:
 				// Fall through to the profile name prompt below.
@@ -306,11 +317,10 @@ a new profile is created.
 		}
 
 		// If no host is available from any source, use the discovery flow
-		// via login.databricks.com.
-		if shouldUseDiscovery(authArguments.Host, args, existingProfile) {
-			if useDeviceCode {
-				return errDeviceCodeNeedsHost
-			}
+		// via login.databricks.com. Discovery opens a browser on this machine,
+		// and device code login is meant for machines without one, so it skips
+		// discovery and setHost resolves the host instead.
+		if !useDeviceCode && shouldUseDiscovery(authArguments.Host, args, existingProfile) {
 			if err := validateDiscoveryFlagCompatibility(cmd); err != nil {
 				return err
 			}
@@ -332,10 +342,16 @@ a new profile is created.
 			return err
 		}
 		// Check the host once it is resolved, before setAccountID prompts for
-		// an account ID for account and unified hosts.
+		// an account ID for account and unified hosts. A workspace without
+		// device code login enabled is rejected later by DeviceCodeChallenge.
 		if useDeviceCode {
-			if err := validateDeviceCodeHost(authArguments); err != nil {
-				return err
+			// The host prompt accepts an empty answer. Without this check,
+			// login would resolve the host from DATABRICKS_HOST instead.
+			if authArguments.Host == "" {
+				return errDeviceCodeNeedsHost
+			}
+			if !isDeviceCodeSupportedForHost(authArguments.Host, authArguments.DiscoveryURL) {
+				return errDeviceCodeWorkspaceOnly
 			}
 		}
 		err = setAccountID(ctx, existingProfile, authArguments)

cmd/auth/login_device_code.go

@@ -6,9 +6,9 @@ import (
 	"fmt"
 	"strings"
 
-	"github.com/databricks/cli/libs/auth"
 	"github.com/databricks/cli/libs/auth/u2m"
 	"github.com/databricks/cli/libs/cmdio"
+	"github.com/databricks/cli/libs/databrickscfg/profile"
 	"github.com/databricks/cli/libs/env"
 	"github.com/spf13/cobra"
 )
@@ -27,14 +27,18 @@ var (
 	errDeviceCodeWorkspaceOnly = errors.New("--use-device-code is only supported for workspace hosts")
 )
 
-// validateDeviceCodeHost returns an error if the host resolved by setHost is
-// an account or unified host. A workspace without device code login enabled
-// is rejected by DeviceCodeChallenge before it creates a device code.
-func validateDeviceCodeHost(authArguments *auth.AuthArguments) error {
-	if isAccountOrUnifiedHost(authArguments.Host, authArguments.DiscoveryURL) {
-		return errDeviceCodeWorkspaceOnly
-	}
-	return nil
+// isDeviceCodeSupportedForHost reports whether host is a workspace host. A
+// unified host is only recognized by its account-scoped discoveryURL, so with
+// an empty discoveryURL only classic account hosts are rejected.
+func isDeviceCodeSupportedForHost(host, discoveryURL string) bool {
+	return !isAccountOrUnifiedHost(host, discoveryURL)
+}
+
+// matchDeviceCodeProfiles matches the profiles that device code login does not
+// always reject. Profiles on unified hosts still match, because the discovery
+// URL that identifies them is not saved on the profile.
+func matchDeviceCodeProfiles(p profile.Profile) bool {
+	return isDeviceCodeSupportedForHost(p.Host, "") && len(u2mResourcesFromProfile(&p)) == 0
 }
 
 // validateDeviceCodeLogin returns an error if device code login is not enabled

cmd/auth/login_device_code_test.go

@@ -5,13 +5,16 @@ import (
 	"context"
 	"errors"
 	"io"
+	"runtime"
 	"strings"
 	"testing"
 	"time"
 
 	"github.com/databricks/cli/libs/auth"
 	"github.com/databricks/cli/libs/auth/u2m"
 	"github.com/databricks/cli/libs/cmdio"
+	"github.com/databricks/cli/libs/cmdio/cmdiotest/termtest"
+	"github.com/databricks/cli/libs/databrickscfg/profile"
 	"github.com/databricks/cli/libs/env"
 	"github.com/databricks/cli/libs/flags"
 	"github.com/spf13/cobra"
@@ -167,7 +170,7 @@ func TestFormatCode(t *testing.T) {
 	assert.Equal(t, "ABC", formatCode("ABC"))
 }
 
-func TestValidateDeviceCodeHost(t *testing.T) {
+func TestIsDeviceCodeSupportedForHost(t *testing.T) {
 	workspace := newDiscoveryServer(t, map[string]any{
 		"account_id":    "acct-123",
 		"workspace_id":  "12345",
@@ -184,23 +187,120 @@ func TestValidateDeviceCodeHost(t *testing.T) {
 	})
 
 	tests := []struct {
-		name    string
-		args    auth.AuthArguments
-		wantErr error
+		name string
+		args auth.AuthArguments
+		want bool
 	}{
-		{name: "workspace host", args: auth.AuthArguments{Host: workspace.URL}},
-		{name: "workspace host with query params", args: auth.AuthArguments{Host: workspace.URL + "?o=12345"}},
-		{name: "classic account host", args: auth.AuthArguments{Host: "https://accounts.cloud.databricks.test"}, wantErr: errDeviceCodeWorkspaceOnly},
-		{name: "unified host", args: auth.AuthArguments{Host: unified.URL}, wantErr: errDeviceCodeWorkspaceOnly},
-		{name: "unified host with query params", args: auth.AuthArguments{Host: unified.URL + "?o=12345"}, wantErr: errDeviceCodeWorkspaceOnly},
-		{name: "unified host with account ID in query params", args: auth.AuthArguments{Host: unifiedWithoutAccountID.URL + "?a=acct-123"}, wantErr: errDeviceCodeWorkspaceOnly},
-		{name: "unified host with account ID flag", args: auth.AuthArguments{Host: unifiedWithoutAccountID.URL, AccountID: "acct-123"}, wantErr: errDeviceCodeWorkspaceOnly},
+		{name: "workspace host", args: auth.AuthArguments{Host: workspace.URL}, want: true},
+		{name: "workspace host with query params", args: auth.AuthArguments{Host: workspace.URL + "?o=12345"}, want: true},
+		{name: "classic account host", args: auth.AuthArguments{Host: "https://accounts.cloud.databricks.test"}},
+		{name: "unified host", args: auth.AuthArguments{Host: unified.URL}},
+		{name: "unified host with query params", args: auth.AuthArguments{Host: unified.URL + "?o=12345"}},
+		{name: "unified host with account ID in query params", args: auth.AuthArguments{Host: unifiedWithoutAccountID.URL + "?a=acct-123"}},
+		{name: "unified host with account ID flag", args: auth.AuthArguments{Host: unifiedWithoutAccountID.URL, AccountID: "acct-123"}},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			args := tt.args
 			require.NoError(t, setHost(t.Context(), nil, &args, nil))
-			assert.ErrorIs(t, validateDeviceCodeHost(&args), tt.wantErr)
+			assert.Equal(t, tt.want, isDeviceCodeSupportedForHost(args.Host, args.DiscoveryURL))
+		})
+	}
+}
+
+// startDeviceCodeLogin runs "auth login --use-device-code" with args in a
+// terminal emulator, using config as the contents of .databrickscfg.
+func startDeviceCodeLogin(t *testing.T, config string, args ...string) *termtest.Term[struct{}] {
+	home := t.TempDir()
+	t.Setenv("HOME", home)
+	t.Setenv("USERPROFILE", home)
+	t.Setenv("DATABRICKS_CONFIG_FILE", writeTempConfig(t, config))
+	t.Setenv("DATABRICKS_AUTH_STORAGE", "plaintext")
+	t.Setenv(deviceCodeEnvVar, "true")
+
+	return termtest.Run(t, func(ctx context.Context) (struct{}, error) {
+		cmd := newLoginCommand(&auth.AuthArguments{})
+		cmd.Flags().String("profile", "", "")
+		cmd.SetContext(ctx)
+		cmd.SetArgs(append([]string{"--use-device-code"}, args...))
+		cmd.SetOut(io.Discard)
+		cmd.SetErr(io.Discard)
+		return struct{}{}, cmd.Execute()
+	})
+}
+
+func TestLoginDeviceCodeNewProfilePromptsForHost(t *testing.T) {
+	tm := startDeviceCodeLogin(t, "", "--profile", "newprof")
+	tm.WaitFor("Databricks host")
+	tm.Type(termtest.KeyEnter)
+	_, err := tm.Result()
+	assert.ErrorIs(t, err, errDeviceCodeNeedsHost)
+}
+
+func TestLoginDeviceCodePickerEmptyHost(t *testing.T) {
+	// termtest feeds stdin through a pipe. On Windows, bubbletea cannot
+	// cancel its read on a pipe, so the picker's reader is still blocked
+	// after it exits and takes the Enter meant for the host prompt.
+	if runtime.GOOS == "windows" {
+		t.Skip("bubbletea cannot cancel reads on non-console input on Windows")
+	}
+	tm := startDeviceCodeLogin(t, "")
+	tm.WaitFor(profilePickerEnterHostLabel)
+	tm.Type(termtest.KeyDown + termtest.KeyEnter)
+	tm.WaitFor("Databricks host")
+	tm.Type(termtest.KeyEnter)
+	_, err := tm.Result()
+	assert.ErrorIs(t, err, errDeviceCodeNeedsHost)
+}
+
+func TestMatchDeviceCodeProfiles(t *testing.T) {
+	tests := []struct {
+		name    string
+		profile profile.Profile
+		want    bool
+	}{
+		{name: "workspace host", profile: profile.Profile{Host: "https://workspace.test"}, want: true},
+		// setHost prompts for the host of a profile without one.
+		{name: "no host", profile: profile.Profile{}, want: true},
+		{name: "classic account host", profile: profile.Profile{Host: "https://accounts.cloud.databricks.test", AccountID: "acct-123"}},
+		{name: "classic account host for DoD", profile: profile.Profile{Host: "https://accounts-dod.cloud.databricks.test", AccountID: "acct-123"}},
+		{
+			name:    "scoped to resources",
+			profile: profile.Profile{Host: "https://workspace.test", AuthType: authTypeDatabricksCLI, Resources: "https://workspace.test/ai-gateway/mcp-services/system.ai.github"},
+		},
+		// Login only reuses the resources of a databricks-cli profile.
+		{
+			name:    "resources on another auth type",
+			profile: profile.Profile{Host: "https://workspace.test", AuthType: "pat", Resources: "https://workspace.test/ai-gateway/mcp-services/system.ai.github"},
+			want:    true,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			assert.Equal(t, tt.want, matchDeviceCodeProfiles(tt.profile))
 		})
 	}
 }
+
+func TestLoginDeviceCodePickerHidesUnsupportedProfiles(t *testing.T) {
+	tm := startDeviceCodeLogin(t, `[workspace-prof]
+host = https://workspace.test
+
+[account-prof]
+host = https://accounts.cloud.databricks.test
+account_id = acct-123
+
+[scoped-prof]
+host = https://workspace.test
+auth_type = databricks-cli
+resources = https://workspace.test/ai-gateway/mcp-services/system.ai.github
+`)
+	tm.WaitFor(profilePickerEnterHostLabel)
+	screen := tm.Snapshot()
+	assert.Contains(t, screen, "workspace-prof")
+	assert.NotContains(t, screen, "account-prof")
+	assert.NotContains(t, screen, "scoped-prof")
+	tm.Type(termtest.KeyCtrlC)
+	_, err := tm.Result()
+	assert.EqualError(t, err, "^C")
+}

docs/device-code-login.md

@@ -33,7 +33,11 @@ Profile my-profile was successfully saved
 ```
 
 1. Run `databricks auth login` with `--use-device-code`. The host can be passed as
-   an argument, with `--host`, or come from an existing profile.
+   an argument, with `--host`, or come from an existing profile. Otherwise the
+   CLI prompts for it. The profile picker leaves out profiles that device code
+   login cannot use: profiles for an account host and profiles saved with
+   `--resource`. Profiles for a unified host are still listed and are rejected
+   after you pick one.
 2. On another device, open the URL that the CLI shows, sign in to the
    workspace if needed, and enter the code. The code expires after about 10
    minutes.