Resolve TokenAudience from default_oidc_audience in host metadata

Signed-off-by: Tanmay Rustagi <tanmay.rustagi@databricks.com>

Author: tanmay-db

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

@@ -931,8 +931,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/oauth/HostMetadata.java

@@ -26,6 +26,9 @@ public class HostMetadata {
   @JsonProperty("host_type")
   private String hostType;
 
+  @JsonProperty("default_oidc_audience")
+  private String defaultOidcAudience;
+
   public HostMetadata() {}
 
   public HostMetadata(String oidcEndpoint, String accountId, String workspaceId) {
@@ -60,4 +63,8 @@ public String getCloud() {
   public String getHostType() {
     return hostType;
   }
+
+  public String getDefaultOidcAudience() {
+    return defaultOidcAudience;
+  }
 }

databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java

@@ -800,6 +800,98 @@ public void testResolveHostMetadataHostTypeUnified() throws IOException {
     }
   }
 
+  // --- resolveHostMetadata default_oidc_audience tests ---
+
+  @Test
+  public void testResolveHostMetadataSetsTokenAudienceFromDefaultOidcAudience() throws IOException {
+    String response =
+        "{\"oidc_endpoint\":\"https://ws.databricks.com/oidc\","
+            + "\"account_id\":\""
+            + DUMMY_ACCOUNT_ID
+            + "\","
+            + "\"default_oidc_audience\":\"https://ws.databricks.com/oidc/v1/token\"}";
+    try (FixtureServer server =
+        new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) {
+      DatabricksConfig config = new DatabricksConfig().setHost(server.getUrl());
+      config.resolve(emptyEnv());
+      config.resolveHostMetadata();
+      assertEquals("https://ws.databricks.com/oidc/v1/token", config.getTokenAudience());
+    }
+  }
+
+  @Test
+  public void testResolveHostMetadataDefaultOidcAudiencePriorityOverAccountIdFallback()
+      throws IOException {
+    // default_oidc_audience should take priority over the account_id fallback for account hosts
+    String response =
+        "{\"oidc_endpoint\":\"https://acc.databricks.com/oidc/accounts/{account_id}\","
+            + "\"account_id\":\""
+            + DUMMY_ACCOUNT_ID
+            + "\","
+            + "\"default_oidc_audience\":\"custom-audience\"}";
+    try (FixtureServer server =
+        new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) {
+      DatabricksConfig config =
+          new DatabricksConfig()
+              .setHost(server.getUrl())
+              .setExperimentalIsUnifiedHost(true)
+              .setAccountId(DUMMY_ACCOUNT_ID);
+      // Note: need two fixtures — resolve() consumes first one (unified host triggers
+      // tryResolveHostMetadata)
+      // Instead, don't set unified — just test direct call
+      config = new DatabricksConfig().setHost(server.getUrl()).setAccountId(DUMMY_ACCOUNT_ID);
+      config.resolve(emptyEnv());
+      config.resolveHostMetadata();
+      // Should use default_oidc_audience, NOT account_id
+      assertEquals("custom-audience", config.getTokenAudience());
+    }
+  }
+
+  @Test
+  public void testResolveHostMetadataDoesNotOverrideExistingTokenAudienceWithOidcAudience()
+      throws IOException {
+    String response =
+        "{\"oidc_endpoint\":\"https://ws.databricks.com/oidc\","
+            + "\"account_id\":\""
+            + DUMMY_ACCOUNT_ID
+            + "\","
+            + "\"default_oidc_audience\":\"metadata-audience\"}";
+    try (FixtureServer server =
+        new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) {
+      DatabricksConfig config =
+          new DatabricksConfig().setHost(server.getUrl()).setTokenAudience("existing-audience");
+      config.resolve(emptyEnv());
+      config.resolveHostMetadata();
+      assertEquals("existing-audience", config.getTokenAudience());
+    }
+  }
+
+  @Test
+  public void testResolveHostMetadataFallsBackToAccountIdWhenNoDefaultOidcAudience()
+      throws IOException {
+    // When no default_oidc_audience, should fall back to account_id for account hosts.
+    // Use unified host flag so getClientType() returns ACCOUNT (no workspaceId).
+    String response =
+        "{\"oidc_endpoint\":\"https://acc.databricks.com/oidc/accounts/{account_id}\","
+            + "\"account_id\":\""
+            + DUMMY_ACCOUNT_ID
+            + "\"}";
+    try (FixtureServer server =
+        new FixtureServer()
+            .with("GET", "/.well-known/databricks-config", response, 200)
+            .with("GET", "/.well-known/databricks-config", response, 200)) {
+      DatabricksConfig config =
+          new DatabricksConfig()
+              .setHost(server.getUrl())
+              .setAccountId(DUMMY_ACCOUNT_ID)
+              .setExperimentalIsUnifiedHost(true);
+      config.resolve(emptyEnv());
+      // resolve() with unified flag triggers tryResolveHostMetadata() consuming first fixture.
+      // Now call again with second fixture to verify account_id fallback.
+      assertEquals(DUMMY_ACCOUNT_ID, config.getTokenAudience());
+    }
+  }
+
   // --- discoveryUrl / OIDC endpoint tests ---
 
   @Test