fix: attach Authorization header to streaming and ApiExecutor requests (#330) (#331)

* fix: Issue #330 added ApiClient#applyAuthHeader as a single entry point for attaching OAuth2 Authorization header.

Includes lazy OAuth2Client initialization on detected CLIENT_CREDENTIALS use.

* fix: #330 use setHeader on request builder to address duplicate Authorization header

* test: #330 improve coverage on ApiClient changeset

* fix: #330 defensively pivot to concurrent map of OAuth2Client to protect multi-tenant credentials

* fix: #330 do not keep raw client secret value in cached keys on heap

* fix: #330 address check-then-act race for access token using single-flight pattern in OAuth2Client

* chore: lint and formatting compliance

* fix: #330 address reviewer feedback for managing ApiClient#applyAuthHeader exception propogation

* fix: #330 consolidate existing AccessToken rather than creating a forked TokenSnapshot

Also improve stream API executor tests per reviewer comments.

* fix: #330 address reviewer input to reorder telemetry call for stability during OAuth2 exchange handling

* fix: #330 address possible redundant token exchange with refactor to strict mutex; includes regression test

Author: cportcvent

src/main/java/dev/openfga/sdk/api/BaseStreamingApi.java

@@ -168,6 +168,8 @@ protected HttpRequest buildHttpRequest(String method, String path, Object body,
             byte[] bodyBytes = objectMapper.writeValueAsBytes(body);
             HttpRequest.Builder requestBuilder = ApiClient.requestBuilder(method, path, bodyBytes, configuration);
 
+            apiClient.applyAuthHeader(requestBuilder, configuration);
+
             // Apply request interceptors if any
             var interceptor = apiClient.getRequestInterceptor();
             if (interceptor != null) {

src/main/java/dev/openfga/sdk/api/OpenFgaApi.java

@@ -15,7 +15,6 @@
 import static dev.openfga.sdk.util.StringUtil.isNullOrWhitespace;
 import static dev.openfga.sdk.util.Validation.assertParamExists;
 
-import dev.openfga.sdk.api.auth.*;
 import dev.openfga.sdk.api.client.*;
 import dev.openfga.sdk.api.configuration.*;
 import dev.openfga.sdk.api.model.BatchCheckRequest;
@@ -69,7 +68,6 @@ public class OpenFgaApi {
     private final Configuration configuration;
 
     private final ApiClient apiClient;
-    private final OAuth2Client oAuth2Client;
     private final Telemetry telemetry;
 
     public OpenFgaApi(Configuration configuration) throws FgaInvalidParameterException {
@@ -89,12 +87,6 @@ public OpenFgaApi(Configuration configuration, ApiClient apiClient, Telemetry te
         this.configuration = configuration;
         this.telemetry = telemetry;
 
-        if (configuration.getCredentials().getCredentialsMethod() == CredentialsMethod.CLIENT_CREDENTIALS) {
-            this.oAuth2Client = new OAuth2Client(configuration, apiClient);
-        } else {
-            this.oAuth2Client = null;
-        }
-
         var defaultHeaders = configuration.getDefaultHeaders();
         if (defaultHeaders != null) {
             apiClient.addRequestInterceptor(httpRequest -> defaultHeaders.forEach(httpRequest::setHeader));
@@ -1294,10 +1286,7 @@ private HttpRequest buildHttpRequestWithPublisher(
         httpRequest.header("Content-Type", "application/json");
         httpRequest.header("Accept", "application/json");
 
-        if (configuration.getCredentials().getCredentialsMethod() != CredentialsMethod.NONE) {
-            String accessToken = getAccessToken(configuration);
-            httpRequest.header("Authorization", "Bearer " + accessToken);
-        }
+        apiClient.applyAuthHeader(httpRequest, configuration);
 
         if (configuration.getUserAgent() != null) {
             httpRequest.header("User-Agent", configuration.getUserAgent());
@@ -1337,29 +1326,4 @@ private String pathWithParams(String basePath, Object... params) {
         }
         return path.toString();
     }
-
-    /**
-     * Get an access token. Expects that configuration is valid (meaning it can
-     * pass {@link Configuration#assertValid()}) and expects that if the
-     * CredentialsMethod is CLIENT_CREDENTIALS that a valid {@link OAuth2Client}
-     * has been initialized. Otherwise, it will throw an IllegalStateException.
-     * @throws IllegalStateException when the configuration is invalid
-     */
-    private String getAccessToken(Configuration configuration) throws ApiException {
-        CredentialsMethod credentialsMethod = configuration.getCredentials().getCredentialsMethod();
-
-        if (credentialsMethod == CredentialsMethod.API_TOKEN) {
-            return configuration.getCredentials().getApiToken().getToken();
-        }
-
-        if (credentialsMethod == CredentialsMethod.CLIENT_CREDENTIALS) {
-            try {
-                return oAuth2Client.getAccessToken().get();
-            } catch (Exception e) {
-                throw new ApiException(e);
-            }
-        }
-
-        throw new IllegalStateException("Configuration is invalid.");
-    }
 }

src/main/java/dev/openfga/sdk/api/auth/AccessToken.java

@@ -5,19 +5,25 @@
 import dev.openfga.sdk.constants.FgaConstants;
 import java.time.Instant;
 import java.time.temporal.ChronoUnit;
-import java.util.Random;
-
-class AccessToken {
+import java.util.concurrent.ThreadLocalRandom;
+
+/**
+ * Immutable snapshot of an access token and its expiry time. The snapshot is valid if the token is non-empty
+ * and the current time is before the expiry time minus a buffer to ensure that callers receive a valid token
+ * even if there is some clock skew or delay between retrieval and use.
+ */
+record AccessToken(String token, Instant expiresAt) {
     private static final int TOKEN_EXPIRY_BUFFER_THRESHOLD_IN_SEC = FgaConstants.TOKEN_EXPIRY_THRESHOLD_BUFFER_IN_SEC;
     // We add some jitter so that token refreshes are less likely to collide
     private static final int TOKEN_EXPIRY_JITTER_IN_SEC = FgaConstants.TOKEN_EXPIRY_JITTER_IN_SEC;
 
-    private Instant expiresAt;
+    static final AccessToken EMPTY = new AccessToken(null, null);
 
-    private final Random random = new Random();
-    private String token;
+    AccessToken {
+        expiresAt = expiresAt != null ? expiresAt.truncatedTo(ChronoUnit.SECONDS) : null;
+    }
 
-    public boolean isValid() {
+    boolean isValid() {
         if (isNullOrWhitespace(token)) {
             return false;
         }
@@ -31,24 +37,9 @@ public boolean isValid() {
         // to account for multiple calls to `isValid` at the same time and prevent multiple refresh calls
         Instant expiresWithLeeway = expiresAt
                 .minusSeconds(TOKEN_EXPIRY_BUFFER_THRESHOLD_IN_SEC)
-                .minusSeconds(random.nextInt(TOKEN_EXPIRY_JITTER_IN_SEC))
+                .minusSeconds(ThreadLocalRandom.current().nextInt(TOKEN_EXPIRY_JITTER_IN_SEC))
                 .truncatedTo(ChronoUnit.SECONDS);
 
         return Instant.now().truncatedTo(ChronoUnit.SECONDS).isBefore(expiresWithLeeway);
     }
-
-    public String getToken() {
-        return token;
-    }
-
-    public void setExpiresAt(Instant expiresAt) {
-        if (expiresAt != null) {
-            // Truncate to seconds to zero out the milliseconds to keep comparison simpler
-            this.expiresAt = expiresAt.truncatedTo(ChronoUnit.SECONDS);
-        }
-    }
-
-    public void setToken(String token) {
-        this.token = token;
-    }
 }

src/main/java/dev/openfga/sdk/api/auth/OAuth2Client.java

@@ -4,20 +4,20 @@
 import dev.openfga.sdk.api.configuration.*;
 import dev.openfga.sdk.errors.ApiException;
 import dev.openfga.sdk.errors.FgaInvalidParameterException;
-import dev.openfga.sdk.telemetry.Attribute;
 import dev.openfga.sdk.telemetry.Telemetry;
 import java.net.URI;
 import java.net.http.HttpRequest;
 import java.time.Instant;
 import java.util.HashMap;
-import java.util.Map;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicReference;
 
 public class OAuth2Client {
     private static final String DEFAULT_API_TOKEN_ISSUER_PATH = "/oauth/token";
 
     private final ApiClient apiClient;
-    private final AccessToken token = new AccessToken();
+    private final AtomicReference<AccessToken> snapshot = new AtomicReference<>(AccessToken.EMPTY);
+    private final AtomicReference<CompletableFuture<String>> inFlight = new AtomicReference<>();
     private final CredentialsFlowRequest authRequest;
     private final Configuration config;
     private final Telemetry telemetry;
@@ -45,26 +45,60 @@ public OAuth2Client(Configuration configuration, ApiClient apiClient) throws Fga
     }
 
     /**
-     * Gets an access token, handling exchange when necessary. The access token is naively cached in memory until it
-     * expires.
+     * Gets an access token, handling exchange when necessary. The token is cached as an immutable
+     * snapshot until it expires. Concurrent calls are deduplicated: only one exchange is in flight
+     * at a time; other callers join the same future rather than issuing redundant requests.
      *
      * @return An access token in a {@link CompletableFuture}
      */
     public CompletableFuture<String> getAccessToken() throws FgaInvalidParameterException, ApiException {
-        if (!token.isValid()) {
-            return exchangeToken().thenCompose(response -> {
-                token.setToken(response.getAccessToken());
-                token.setExpiresAt(Instant.now().plusSeconds(response.getExpiresInSeconds()));
-
-                Map<Attribute, String> attributesMap = new HashMap<>();
-
-                telemetry.metrics().credentialsRequest(1L, attributesMap);
-
-                return CompletableFuture.completedFuture(token.getToken());
-            });
+        // Fast path (lock-free): return cached token if still valid.
+        AccessToken current = snapshot.get();
+        if (current.isValid()) {
+            return CompletableFuture.completedFuture(current.token());
         }
 
-        return CompletableFuture.completedFuture(token.getToken());
+        // Slow path: decide under the lock who starts the exchange.
+        synchronized (this) {
+            // Double-check: another thread may have refreshed while we waited.
+            AccessToken rechecked = snapshot.get();
+            if (rechecked.isValid()) {
+                return CompletableFuture.completedFuture(rechecked.token());
+            }
+
+            // Join an existing in-flight exchange.
+            CompletableFuture<String> existing = inFlight.get();
+            if (existing != null) {
+                return existing;
+            }
+
+            // Start a new exchange and publish the future so other callers join it.
+            CompletableFuture<String> promise = new CompletableFuture<>();
+            inFlight.set(promise);
+
+            try {
+                exchangeToken().whenComplete((response, ex) -> {
+                    if (ex != null) {
+                        inFlight.set(null);
+                        promise.completeExceptionally(ex);
+                    } else {
+                        String token = response.getAccessToken();
+                        // Write snapshot before clearing the gate so any new caller that arrives
+                        // after inFlight becomes null immediately sees a valid token.
+                        snapshot.set(new AccessToken(token, Instant.now().plusSeconds(response.getExpiresInSeconds())));
+                        inFlight.set(null);
+                        promise.complete(token);
+                        telemetry.metrics().credentialsRequest(1L, new HashMap<>());
+                    }
+                });
+            } catch (Exception e) {
+                inFlight.set(null);
+                promise.completeExceptionally(e);
+                throw e;
+            }
+
+            return promise;
+        }
     }
 
     /**

src/main/java/dev/openfga/sdk/api/client/ApiClient.java

@@ -7,7 +7,12 @@
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.SerializationFeature;
 import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule;
+import dev.openfga.sdk.api.auth.OAuth2Client;
+import dev.openfga.sdk.api.configuration.ClientCredentials;
 import dev.openfga.sdk.api.configuration.Configuration;
+import dev.openfga.sdk.api.configuration.Credentials;
+import dev.openfga.sdk.api.configuration.CredentialsMethod;
+import dev.openfga.sdk.errors.ApiException;
 import dev.openfga.sdk.errors.FgaInvalidParameterException;
 import dev.openfga.sdk.util.StringUtil;
 import java.io.InputStream;
@@ -16,7 +21,14 @@
 import java.net.http.HttpClient;
 import java.net.http.HttpRequest;
 import java.net.http.HttpResponse;
+import java.security.MessageDigest;
+import java.security.NoSuchAlgorithmException;
 import java.time.Duration;
+import java.util.Arrays;
+import java.util.Objects;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.ExecutionException;
 import java.util.function.Consumer;
 import org.openapitools.jackson.nullable.JsonNullableModule;
 
@@ -41,6 +53,7 @@ public class ApiClient {
     private Consumer<HttpRequest.Builder> interceptor;
     private Consumer<HttpResponse<InputStream>> responseInterceptor;
     private Consumer<HttpResponse<String>> asyncResponseInterceptor;
+    private final ConcurrentMap<CredentialsCacheKey, OAuth2Client> oAuth2Clients = new ConcurrentHashMap<>();
 
     /**
      * Create an instance of ApiClient.
@@ -324,4 +337,118 @@ public ApiClient setAsyncResponseInterceptor(Consumer<HttpResponse<String>> inte
     public Consumer<HttpResponse<String>> getAsyncResponseInterceptor() {
         return asyncResponseInterceptor;
     }
+
+    /**
+     * Applies the {@code Authorization: Bearer <token>} header to the request builder based on the
+     * supplied configuration's {@link Credentials}. This is the single entry point for attaching
+     * auth to outbound requests across the SDK — every request builder should delegate here.
+     *
+     * <ul>
+     *   <li>{@link CredentialsMethod#NONE}: no header is applied.</li>
+     *   <li>{@link CredentialsMethod#API_TOKEN}: the static API token from the configuration is used.</li>
+     *   <li>{@link CredentialsMethod#CLIENT_CREDENTIALS}: an {@link OAuth2Client} performs the
+     *       client-credentials exchange and caches the token on this {@code ApiClient} until expiry.
+     *       The client is lazily created from {@code configuration} on first use.</li>
+     * </ul>
+     *
+     * @param requestBuilder the request builder to mutate.
+     * @param configuration  the configuration that supplies credentials.
+     * @throws ApiException                 if CLIENT_CREDENTIALS token exchange fails.
+     * @throws FgaInvalidParameterException if the configuration is invalid when lazily creating
+     *                                      an {@link OAuth2Client}.
+     */
+    public void applyAuthHeader(HttpRequest.Builder requestBuilder, Configuration configuration)
+            throws ApiException, FgaInvalidParameterException {
+
+        Credentials credentials = configuration.getCredentials();
+        if (credentials == null) {
+            return;
+        }
+
+        CredentialsMethod method = credentials.getCredentialsMethod();
+        if (method == null || method == CredentialsMethod.NONE) {
+            return;
+        }
+
+        String accessToken;
+        switch (method) {
+            case API_TOKEN:
+                accessToken = credentials.getApiToken().getToken();
+                break;
+            case CLIENT_CREDENTIALS:
+                try {
+                    accessToken =
+                            ensureOAuth2Client(configuration).getAccessToken().get();
+                } catch (InterruptedException e) {
+                    Thread.currentThread().interrupt();
+                    throw new ApiException(e);
+                } catch (ExecutionException e) {
+                    Throwable cause = e.getCause();
+                    if (cause instanceof ApiException) {
+                        throw (ApiException) cause;
+                    }
+                    throw new ApiException(cause != null ? cause : e);
+                }
+                break;
+            default:
+                throw new IllegalStateException("Unknown credentials method: " + method);
+        }
+
+        requestBuilder.setHeader("Authorization", "Bearer " + accessToken);
+    }
+
+    private OAuth2Client ensureOAuth2Client(Configuration configuration) throws FgaInvalidParameterException {
+        ClientCredentials cc = configuration.getCredentials().getClientCredentials();
+        CredentialsCacheKey key = new CredentialsCacheKey(cc);
+        OAuth2Client existing = oAuth2Clients.get(key);
+        if (existing != null) {
+            return existing;
+        }
+        OAuth2Client created = new OAuth2Client(configuration, this);
+        OAuth2Client prior = oAuth2Clients.putIfAbsent(key, created);
+        return prior != null ? prior : created;
+    }
+
+    private static final class CredentialsCacheKey {
+        private final String clientId;
+        private final byte[] clientSecretHash;
+        private final String apiTokenIssuer;
+        private final String apiAudience;
+        private final String scopes;
+
+        CredentialsCacheKey(ClientCredentials cc) {
+            this.clientId = cc.getClientId();
+            this.clientSecretHash = sha256(cc.getClientSecret());
+            this.apiTokenIssuer = cc.getApiTokenIssuer();
+            this.apiAudience = cc.getApiAudience();
+            this.scopes = cc.getScopes();
+        }
+
+        private static byte[] sha256(String value) {
+            try {
+                return MessageDigest.getInstance("SHA-256").digest(value == null ? new byte[0] : value.getBytes(UTF_8));
+            } catch (NoSuchAlgorithmException e) {
+                throw new IllegalStateException("SHA-256 not available", e);
+            }
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (!(o instanceof CredentialsCacheKey)) return false;
+            CredentialsCacheKey that = (CredentialsCacheKey) o;
+            return Objects.equals(clientId, that.clientId)
+                    && Arrays.equals(clientSecretHash, that.clientSecretHash)
+                    && Objects.equals(apiTokenIssuer, that.apiTokenIssuer)
+                    && Objects.equals(apiAudience, that.apiAudience)
+                    && Objects.equals(scopes, that.scopes);
+        }
+
+        @Override
+        public int hashCode() {
+            int result = Objects.hash(clientId, apiTokenIssuer, apiAudience, scopes);
+            result = 31 * result + Arrays.hashCode(clientSecretHash);
+            return result;
+        }
+    }
 }