Remove unified flag usage, rely on host metadata

account_functions.go

@@ -63,10 +63,6 @@ func (c *AccountClient) GetWorkspaceClient(ws provisioning.Workspace) (*Workspac
 	}
 	cfg.AzureResourceID = ws.AzureResourceId()
 	cfg.WorkspaceID = fmt.Sprintf("%d", ws.WorkspaceId)
-	if c.Config.Experimental_IsUnifiedHost {
-		cfg.AccountID = c.Config.AccountID
-		cfg.Experimental_IsUnifiedHost = true
-	}
 	w, err := NewWorkspaceClient((*Config)(cfg))
 	if err != nil {
 		return nil, err

client/client_test.go

@@ -22,6 +22,15 @@ import (
 type hc func(r *http.Request) (*http.Response, error)
 
 func (cb hc) RoundTrip(r *http.Request) (*http.Response, error) {
+	// Return 404 for host metadata endpoint to prevent config resolution
+	// from interfering with test assertions.
+	if r.Method == "GET" && r.URL.Path == "/.well-known/databricks-config" {
+		return &http.Response{
+			StatusCode: 404,
+			Body:       io.NopCloser(strings.NewReader(`{"error_code":"NOT_FOUND","message":"host metadata endpoint auto-stubbed by test framework"}`)),
+			Request:    r,
+		}, nil
+	}
 	return cb(r)
 }
 

config/auth_azure_msi.go

@@ -14,8 +14,10 @@ import (
 	"golang.org/x/oauth2"
 )
 
-var errInvalidToken = errors.New("invalid token")
-var errInvalidTokenExpiry = errors.New("invalid token expiry")
+var (
+	errInvalidToken       = errors.New("invalid token")
+	errInvalidTokenExpiry = errors.New("invalid token expiry")
+)
 
 // well-known URL for Azure Instance Metadata Service (IMDS)
 // https://learn.microsoft.com/en-us/azure-stack/user/instance-metadata-service
@@ -24,19 +26,20 @@ var instanceMetadataPrefix = "http://169.254.169.254/metadata"
 // timeout to wait for IMDS response
 const azureMsiTimeout = 10 * time.Second
 
-type AzureMsiCredentials struct {
-}
+type AzureMsiCredentials struct{}
 
 func (c AzureMsiCredentials) Name() string {
 	return "azure-msi"
 }
 
 func (c AzureMsiCredentials) Configure(ctx context.Context, cfg *Config) (credentials.CredentialsProvider, error) {
-	if !cfg.IsAzure() || !cfg.AzureUseMSI || (cfg.AzureResourceID == "" && cfg.ConfigType() == WorkspaceConfig) {
+	if !cfg.IsAzure() || !cfg.AzureUseMSI || (cfg.AzureResourceID == "" && cfg.Host == "") {
 		return nil, nil
 	}
 	env := cfg.Environment()
-	if !cfg.IsAccountClient() {
+	// If the host is not set, we need to resolve it from the Azure Resource ID.
+	// This is only needed for Workspaces, because Accounts always have a host.
+	if cfg.Host == "" {
 		err := cfg.azureEnsureWorkspaceUrl(ctx, c)
 		if err != nil {
 			return nil, fmt.Errorf("resolve host: %w", err)

config/auth_azure_msi_test.go

@@ -121,6 +121,33 @@ func TestMsiTokenNotFound(t *testing.T) {
 	require.ErrorIs(t, err, apierr.ErrNotFound)
 }
 
+func TestMsiHappyFlowWithHostAndNoResourceID(t *testing.T) {
+	assertHeaders(t, &Config{
+		Host:        "https://adb-123.4.azuredatabricks.net",
+		AzureUseMSI: true,
+		AuthType:    "azure-msi",
+		HTTPTransport: fixtures.MappingTransport{
+			"GET /metadata/identity/oauth2/token?api-version=2018-02-01&resource=2ff814a6-3304-4ab8-85cb-cd0e6f879c1d": {
+				ExpectedHeaders: map[string]string{
+					"Accept":   "application/json",
+					"Metadata": "true",
+				},
+				Response: someValidToken("cde"),
+			},
+			"GET /metadata/identity/oauth2/token?api-version=2018-02-01&resource=https%3A%2F%2Fmanagement.core.windows.net%2F": {
+				ExpectedHeaders: map[string]string{
+					"Accept":   "application/json",
+					"Metadata": "true",
+				},
+				Response: someValidToken("def"),
+			},
+		},
+	}, map[string]string{
+		"Authorization":                          "Bearer cde",
+		"X-Databricks-Azure-Sp-Management-Token": "def",
+	})
+}
+
 func TestMsiInvalidTokenExpiry(t *testing.T) {
 	_, err := authenticateRequest(&Config{
 		AzureUseMSI:     true,

config/auth_oidc_test.go

@@ -74,6 +74,9 @@ func TestGithubOIDC_Scopes(t *testing.T) {
 						"expires_in":   3600,
 					})
 
+				case "/.well-known/databricks-config":
+					http.Error(w, "Not found", http.StatusNotFound)
+
 				default:
 					t.Errorf("Unexpected request: %s %s", r.Method, r.URL.Path)
 					http.Error(w, "Not found", http.StatusNotFound)

config/cli_token_source.go

@@ -70,14 +70,6 @@ func buildCliCommands(cliPath string, cfg *Config) (primaryCmd []string, hostCmd
 func buildHostCommand(cliPath string, cfg *Config) []string {
 	cmd := []string{cliPath, "auth", "token", "--host", cfg.Host}
 	switch cfg.HostType() {
-	case UnifiedHost:
-		cmd = append(cmd, "--experimental-is-unified-host")
-		if cfg.AccountID != "" {
-			cmd = append(cmd, "--account-id", cfg.AccountID)
-		}
-		if cfg.WorkspaceID != "" {
-			cmd = append(cmd, "--workspace-id", cfg.WorkspaceID)
-		}
 	case AccountHost:
 		cmd = append(cmd, "--account-id", cfg.AccountID)
 	}

config/cli_token_source_test.go

@@ -152,50 +152,13 @@ func TestBuildCliCommands(t *testing.T) {
 			wantCmd: []string{cliPath, "auth", "token", "--host", accountHost, "--account-id", accountID},
 		},
 		{
-			name: "unified host with account ID and workspace ID",
+			name: "former unified host treated as workspace",
 			cfg: &Config{
-				Host:                       unifiedHost,
-				Experimental_IsUnifiedHost: true,
-				AccountID:                  accountID,
-				WorkspaceID:                workspaceID,
+				Host:        unifiedHost,
+				AccountID:   accountID,
+				WorkspaceID: workspaceID,
 			},
-			wantCmd: []string{cliPath, "auth", "token", "--host", unifiedHost, "--experimental-is-unified-host", "--account-id", accountID, "--workspace-id", workspaceID},
-		},
-		{
-			name: "unified host with account ID only",
-			cfg: &Config{
-				Host:                       unifiedHost,
-				Experimental_IsUnifiedHost: true,
-				AccountID:                  accountID,
-			},
-			wantCmd: []string{cliPath, "auth", "token", "--host", unifiedHost, "--experimental-is-unified-host", "--account-id", accountID},
-		},
-		{
-			name: "unified host with workspace ID only",
-			cfg: &Config{
-				Host:                       unifiedHost,
-				Experimental_IsUnifiedHost: true,
-				WorkspaceID:                workspaceID,
-			},
-			wantCmd: []string{cliPath, "auth", "token", "--host", unifiedHost, "--experimental-is-unified-host", "--workspace-id", workspaceID},
-		},
-		{
-			name: "unified host with no account ID or workspace ID",
-			cfg: &Config{
-				Host:                       unifiedHost,
-				Experimental_IsUnifiedHost: true,
-			},
-			wantCmd: []string{cliPath, "auth", "token", "--host", unifiedHost, "--experimental-is-unified-host"},
-		},
-		{
-			// Explicit false should fall back to the standard host type detection
-			name: "unified host false with account host",
-			cfg: &Config{
-				Host:                       accountHost,
-				Experimental_IsUnifiedHost: false,
-				AccountID:                  accountID,
-			},
-			wantCmd: []string{cliPath, "auth", "token", "--host", accountHost, "--account-id", accountID},
+			wantCmd: []string{cliPath, "auth", "token", "--host", unifiedHost},
 		},
 		{
 			name:        "profile uses --profile with --host fallback",

config/config.go

@@ -379,14 +379,9 @@ func (c *Config) IsAws() bool {
 }
 
 // IsAccountClient returns true if client is configured for Accounts API.
-// Panics if the config has the unified host flag set.
 //
 // Deprecated: Use HostType() if possible, or ConfigType() if necessary.
 func (c *Config) IsAccountClient() bool {
-	if c.Experimental_IsUnifiedHost {
-		panic("IsAccountClient cannot be used with unified hosts; use HostType() instead")
-	}
-
 	if c.AccountID != "" && c.isTesting {
 		return true
 	}
@@ -413,11 +408,13 @@ func normalizedHost(host string) string {
 }
 
 // HostType returns the type of host that the client is configured for.
+// HostType now only returns WorkspaceHost or AccountHost. UnifiedHost is
+// deprecated; host metadata resolution handles unified host behavior.
 func (c *Config) HostType() HostType {
+	// TODO: Remove this after TF updates its code.
 	if c.Experimental_IsUnifiedHost {
 		return UnifiedHost
 	}
-
 	// TODO: Refactor tests so that this is not needed.
 	if c.AccountID != "" && c.isTesting {
 		return AccountHost
@@ -452,15 +449,6 @@ func (c *Config) ConfigType() ConfigType {
 		return AccountConfig
 	case WorkspaceHost:
 		return WorkspaceConfig
-	case UnifiedHost:
-		if c.AccountID == "" {
-			// All unified host configs must have an account ID
-			return InvalidConfig
-		}
-		if c.WorkspaceID != "" {
-			return WorkspaceConfig
-		}
-		return AccountConfig
 	default:
 		return InvalidConfig
 	}
@@ -520,9 +508,7 @@ func (c *Config) EnsureResolved() error {
 	}
 	slices.Sort(c.Scopes)
 	c.Scopes = slices.Compact(c.Scopes)
-	if c.Experimental_IsUnifiedHost {
-		c.resolveHostMetadata(ctx)
-	}
+	c.resolveHostMetadata(ctx)
 	c.resolved = true
 	return nil
 }
@@ -653,8 +639,6 @@ func (c *Config) getOidcEndpoints(ctx context.Context) (*u2m.OAuthAuthorizationS
 	switch c.HostType() {
 	case AccountHost:
 		return oauthClient.GetAccountOAuthEndpoints(ctx, host, c.AccountID)
-	case UnifiedHost:
-		return oauthClient.GetUnifiedOAuthEndpoints(ctx, host, c.AccountID)
 	case WorkspaceHost:
 		return oauthClient.GetWorkspaceOAuthEndpoints(ctx, host)
 	default:
@@ -733,8 +717,6 @@ func (c *Config) getOAuthArgument() (u2m.OAuthArgument, error) {
 	switch c.HostType() {
 	case AccountHost:
 		return u2m.NewProfileAccountOAuthArgument(host, c.AccountID, profile)
-	case UnifiedHost:
-		return u2m.NewProfileUnifiedOAuthArgument(host, c.AccountID, profile)
 	case WorkspaceHost:
 		return u2m.NewProfileWorkspaceOAuthArgument(host, profile)
 	default: