From 2dd4a6da83dd5de9f279c0b2bfe37d3abf7a74a8 Mon Sep 17 00:00:00 2001
From: Hector Castejon Diaz <hector.castejon@databricks.com>
Date: Thu, 19 Mar 2026 09:20:09 +0000
Subject: [PATCH] Resolve TokenAudience from host metadata for account hosts

Port of Go SDK #1543. When resolveHostMetadata() runs on an account
host and tokenAudience is not already set, automatically sets it to
the accountId. This enables OIDC token exchange to work correctly
for account-level operations without explicit TOKEN_AUDIENCE config.

Co-authored-by: Isaac
---
 .../databricks/sdk/core/DatabricksConfig.java |  4 ++
 .../sdk/core/DatabricksConfigTest.java        | 46 +++++++++++++++++++
 2 files changed, 50 insertions(+)

diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java
index b2b0d75e4..43c1f5136 100644
--- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java
+++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java
@@ -907,6 +907,10 @@ 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 && getClientType() == ClientType.ACCOUNT && accountId != null) {
+      tokenAudience = accountId;
+    }
   }
 
   private OpenIDConnectEndpoints fetchOidcEndpointsFromDiscovery() {
diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java
index c6e149d76..0e86b4ef4 100644
--- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java
+++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java
@@ -586,6 +586,52 @@ public void testResolveHostMetadataRaisesOnHttpError() throws IOException {
     }
   }
 
+  @Test
+  public void testResolveHostMetadataSetsTokenAudienceForAccountHost() throws IOException {
+    // For a unified host with no workspaceId (ACCOUNT client type), resolveHostMetadata should
+    // set tokenAudience to accountId when not already configured.
+    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())
+              .setExperimentalIsUnifiedHost(true)
+              .setAccountId(DUMMY_ACCOUNT_ID);
+      config.resolve(emptyEnv());
+      // Client type should be ACCOUNT (unified host, no workspaceId)
+      assertEquals(ClientType.ACCOUNT, config.getClientType());
+      config.resolveHostMetadata();
+      assertEquals(DUMMY_ACCOUNT_ID, config.getTokenAudience());
+    }
+  }
+
+  @Test
+  public void testResolveHostMetadataDoesNotOverwriteTokenAudience() throws IOException {
+    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)) {
+      DatabricksConfig config =
+          new DatabricksConfig()
+              .setHost(server.getUrl())
+              .setAccountId(DUMMY_ACCOUNT_ID)
+              .setTokenAudience("custom-audience");
+      config.resolve(emptyEnv());
+      config.resolveHostMetadata();
+      assertEquals("custom-audience", config.getTokenAudience());
+    }
+  }
+
   // --- tryResolveHostMetadata (config init) tests ---
 
   @Test