[Internal] Use resolved host type from host metadata in getHostType()

databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java

@@ -162,6 +162,12 @@ public class DatabricksConfig {
 
   private DatabricksEnvironment databricksEnvironment;
 
+  /**
+   * The host type resolved from the /.well-known/databricks-config discovery endpoint. When set,
+   * this takes priority over URL-based host type detection in {@link #getHostType()}.
+   */
+  private HostType resolvedHostType;
+
   /**
    * When using Workload Identity Federation, the audience to specify when fetching an ID token from
    * the ID token supplier.
@@ -714,6 +720,17 @@ public DatabricksConfig setDisableOauthRefreshToken(boolean disable) {
     return this;
   }
 
+  /** Returns the host type resolved from host metadata, or {@code null} if not yet resolved. */
+  HostType getResolvedHostType() {
+    return resolvedHostType;
+  }
+
+  /** Sets the resolved host type. Package-private for testing. */
+  DatabricksConfig setResolvedHostType(HostType resolvedHostType) {
+    this.resolvedHostType = resolvedHostType;
+    return this;
+  }
+
   public boolean isAzure() {
     if (azureWorkspaceResourceId != null) {
       return true;
@@ -743,8 +760,15 @@ public boolean isAccountClient() {
     return host.startsWith("https://accounts.") || host.startsWith("https://accounts-dod.");
   }
 
-  /** Returns the host type based on the host URL pattern. */
+  /**
+   * Returns the host type based on configuration settings and host URL. When host metadata has been
+   * resolved (via /.well-known/databricks-config), the resolved host type is returned directly.
+   * Otherwise, the host type is inferred from URL patterns as a fallback.
+   */
   public HostType getHostType() {
+    if (resolvedHostType != null) {
+      return resolvedHostType;
+    }
     if (host == null) {
       return HostType.WORKSPACE;
     }
@@ -866,6 +890,13 @@ void resolveHostMetadata() throws IOException {
       LOG.debug("Resolved workspace_id from host metadata: \"{}\"", meta.getWorkspaceId());
       workspaceId = meta.getWorkspaceId();
     }
+    if (resolvedHostType == null && meta.getHostType() != null) {
+      HostType ht = HostType.fromApiValue(meta.getHostType());
+      if (ht != null) {
+        LOG.debug("Resolved host_type from host metadata: \"{}\"", ht);
+        resolvedHostType = ht;
+      }
+    }
     if (discoveryUrl == null) {
       if (meta.getOidcEndpoint() == null || meta.getOidcEndpoint().isEmpty()) {
         LOG.warn("Host metadata missing oidc_endpoint; skipping discovery URL resolution");
@@ -884,8 +915,17 @@ void resolveHostMetadata() throws IOException {
       discoveryUrl = oidcUri.resolve(".well-known/oauth-authorization-server").toString();
       LOG.debug("Resolved discovery_url from host metadata: \"{}\"", discoveryUrl);
     }
-    // For account hosts, use the accountId as the token audience if not already set.
+    if (tokenAudience == null
+        && meta.getDefaultOidcAudience() != null
+        && !meta.getDefaultOidcAudience().isEmpty()) {
+      LOG.debug(
+          "Resolved token_audience from host metadata default_oidc_audience: \"{}\"",
+          meta.getDefaultOidcAudience());
+      tokenAudience = meta.getDefaultOidcAudience();
+    }
+    // Fallback: for account hosts, use the accountId as the token audience if not already set.
     if (tokenAudience == null && getClientType() == ClientType.ACCOUNT && accountId != null) {
+      LOG.debug("Setting token_audience to account_id for account host: \"{}\"", accountId);
       tokenAudience = accountId;
     }
   }

databricks-sdk-java/src/main/java/com/databricks/sdk/core/HostType.java

@@ -9,5 +9,28 @@ public enum HostType {
   WORKSPACE,
 
   /** Traditional accounts host. */
-  ACCOUNTS
+  ACCOUNTS,
+
+  /** Unified host supporting both workspace and account operations. */
+  UNIFIED;
+
+  /**
+   * Converts an API-level host type string (e.g. "workspace", "account", "unified") to the
+   * corresponding enum value. Returns {@code null} for unknown or empty values.
+   */
+  public static HostType fromApiValue(String value) {
+    if (value == null || value.isEmpty()) {
+      return null;
+    }
+    switch (value.toLowerCase()) {
+      case "workspace":
+        return WORKSPACE;
+      case "account":
+        return ACCOUNTS;
+      case "unified":
+        return UNIFIED;
+      default:
+        return null;
+    }
+  }
 }

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/HostMetadata.java

@@ -23,6 +23,12 @@ public class HostMetadata {
   @JsonProperty("cloud")
   private String cloud;
 
+  @JsonProperty("host_type")
+  private String hostType;
+
+  @JsonProperty("default_oidc_audience")
+  private String defaultOidcAudience;
+
   public HostMetadata() {}
 
   public HostMetadata(String oidcEndpoint, String accountId, String workspaceId) {
@@ -53,4 +59,12 @@ public String getWorkspaceId() {
   public String getCloud() {
     return cloud;
   }
+
+  public String getHostType() {
+    return hostType;
+  }
+
+  public String getDefaultOidcAudience() {
+    return defaultOidcAudience;
+  }
 }

databricks-sdk-java/src/test/java/com/databricks/sdk/AccountClientTest.java

@@ -3,8 +3,10 @@
 import static org.junit.jupiter.api.Assertions.*;
 
 import com.databricks.sdk.core.DatabricksConfig;
+import com.databricks.sdk.core.FixtureServer;
 import com.databricks.sdk.core.HostType;
 import com.databricks.sdk.service.provisioning.Workspace;
+import java.io.IOException;
 import org.junit.jupiter.api.Test;
 
 public class AccountClientTest {
@@ -31,30 +33,33 @@ public void testGetWorkspaceClientForTraditionalAccount() {
   }
 
   @Test
-  public void testGetWorkspaceClientForUnifiedHost() {
-    String unifiedHost = "https://unified.databricks.com";
-    DatabricksConfig accountConfig =
-        new DatabricksConfig()
-            .setHost(unifiedHost)
-            .setAccountId("test-account")
-            .setToken("test-token");
+  public void testGetWorkspaceClientForUnifiedHost() throws IOException {
+    try (FixtureServer server = new FixtureServer()) {
+      String unifiedHost = server.getUrl();
+      DatabricksConfig accountConfig =
+          new DatabricksConfig()
+              .setHost(unifiedHost)
+              .setAccountId("test-account")
+              .setToken("test-token");
 
-    AccountClient accountClient = new AccountClient(accountConfig);
+      AccountClient accountClient = new AccountClient(accountConfig);
 
-    Workspace workspace = new Workspace();
-    workspace.setWorkspaceId(123456L);
-    workspace.setDeploymentName("test-workspace");
+      Workspace workspace = new Workspace();
+      workspace.setWorkspaceId(123456L);
+      workspace.setDeploymentName("test-workspace");
 
-    WorkspaceClient workspaceClient = accountClient.getWorkspaceClient(workspace);
+      WorkspaceClient workspaceClient = accountClient.getWorkspaceClient(workspace);
 
-    // Should have the same host (unified hosts reuse the same host)
-    assertEquals(unifiedHost, workspaceClient.config().getHost());
+      // Should have the same host (non-matching DNS zone means SPOG path)
+      assertEquals(unifiedHost, workspaceClient.config().getHost());
 
-    // Should have workspace ID set
-    assertEquals("123456", workspaceClient.config().getWorkspaceId());
+      // Should have workspace ID set
+      assertEquals("123456", workspaceClient.config().getWorkspaceId());
 
-    // Host type is WORKSPACE (determined from URL pattern, not unified flag)
-    assertEquals(HostType.WORKSPACE, workspaceClient.config().getHostType());
+      // Host type is WORKSPACE (no resolved host type from metadata, URL doesn't match accounts
+      // pattern)
+      assertEquals(HostType.WORKSPACE, workspaceClient.config().getHostType());
+    }
   }
 
   @Test