Skip to content

Commit 7d4a83a

Browse files
jaycee-licopybara-github
jaycee-li
authored and
copybara-github
committed
fix: refresh the access token on every HTTP attempt so retries and queued async calls do not send an expired token
PiperOrigin-RevId: 1009949015
1 parent 79d09a6 commit 7d4a83a

2 files changed

Lines changed: 249 additions & 2 deletions

File tree

‎src/main/java/com/google/genai/ApiClient.java‎

Lines changed: 41 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -47,10 +47,12 @@
4747
import java.util.stream.Stream;
4848
import okhttp3.Credentials;
4949
import okhttp3.Dispatcher;
50+
import okhttp3.Interceptor;
5051
import okhttp3.MediaType;
5152
import okhttp3.OkHttpClient;
5253
import okhttp3.Request;
5354
import okhttp3.RequestBody;
55+
import okhttp3.Response;
5456
import org.jspecify.annotations.Nullable;
5557

5658
/** Interface for an API client which issues HTTP requests to the GenAI APIs. */
@@ -305,6 +307,8 @@ private OkHttpClient createHttpClient(
305307
HttpRetryOptions retryOptions =
306308
httpOptions.retryOptions().orElse(HttpRetryOptions.builder().build());
307309
builder.addInterceptor(new RetryInterceptor(retryOptions));
310+
// Must come after RetryInterceptor so that it runs once per attempt.
311+
builder.addInterceptor(this::reauthorize);
308312

309313
return builder.build();
310314
}
@@ -525,8 +529,9 @@ private void setHeaders(Request.Builder request, HttpOptions requestHttpOptions)
525529
} catch (IOException e) {
526530
throw new GenAiIOException("Failed to refresh credentials.", e);
527531
}
528-
String accessToken = cred.getAccessToken().getTokenValue();
529-
request.header("Authorization", "Bearer " + accessToken);
532+
String authorization = "Bearer " + cred.getAccessToken().getTokenValue();
533+
request.header("Authorization", authorization);
534+
request.tag(SdkAuthorization.class, new SdkAuthorization(authorization));
530535

531536
if (cred.getQuotaProjectId() != null) {
532537
request.header("x-goog-user-project", cred.getQuotaProjectId());
@@ -536,6 +541,40 @@ private void setHeaders(Request.Builder request, HttpOptions requestHttpOptions)
536541
}
537542
}
538543

544+
/**
545+
* Re-applies credentials before each attempt, so that a retried or queued request does not
546+
* carry an access token that expired after the request was built.
547+
*
548+
* <p>The refresh goes through {@link #setHeaders} for the same Boq conformance reason as {@link
549+
* #refreshAndGetAccessToken}. A request whose Authorization header was changed after the SDK set
550+
* it is sent unchanged.
551+
*/
552+
private Response reauthorize(Interceptor.Chain chain) throws IOException {
553+
Request request = chain.request();
554+
SdkAuthorization sdkAuthorization = request.tag(SdkAuthorization.class);
555+
if (sdkAuthorization == null
556+
|| !sdkAuthorization.headerValue.equals(request.header("Authorization"))) {
557+
return chain.proceed(request);
558+
}
559+
Request.Builder builder = request.newBuilder();
560+
try {
561+
setHeaders(builder, HttpOptions.builder().build());
562+
} catch (GenAiIOException e) {
563+
// OkHttp only reports IOExceptions to async callbacks; others kill the dispatcher thread.
564+
throw new IOException(e.getMessage(), e);
565+
}
566+
return chain.proceed(builder.build());
567+
}
568+
569+
/** The Authorization header value that {@link #setHeaders} set on a request. */
570+
private static final class SdkAuthorization {
571+
final String headerValue;
572+
573+
SdkAuthorization(String headerValue) {
574+
this.headerValue = headerValue;
575+
}
576+
}
577+
539578
/**
540579
* Refreshes credentials (if needed) and returns the current access token.
541580
*

‎src/test/java/com/google/genai/HttpApiClientTest.java‎

Lines changed: 208 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,13 +41,21 @@
4141
import java.net.InetSocketAddress;
4242
import java.net.Proxy;
4343
import java.time.Duration;
44+
import java.util.ArrayList;
45+
import java.util.Collections;
46+
import java.util.Date;
47+
import java.util.List;
4448
import java.util.Map;
4549
import java.util.Optional;
4650
import java.util.concurrent.CompletableFuture;
51+
import java.util.concurrent.CountDownLatch;
4752
import java.util.concurrent.ExecutionException;
53+
import java.util.concurrent.TimeUnit;
54+
import java.util.concurrent.atomic.AtomicReference;
4855
import okhttp3.Call;
4956
import okhttp3.Callback;
5057
import okhttp3.Dispatcher;
58+
import okhttp3.Interceptor;
5159
import okhttp3.OkHttpClient;
5260
import okhttp3.Protocol;
5361
import okhttp3.Request;
@@ -92,6 +100,11 @@ public class HttpApiClientTest {
92100
GoogleCredentials.newBuilder()
93101
.setAccessToken(AccessToken.newBuilder().setTokenValue("").build())
94102
.build();
103+
private static final HttpOptions FAST_RETRY_HTTP_OPTIONS =
104+
HttpOptions.builder()
105+
.retryOptions(
106+
HttpRetryOptions.builder().attempts(2).initialDelay(0.0).jitter(0.0).build())
107+
.build();
95108

96109
@Mock OkHttpClient mockHttpClient;
97110
@Mock Call mockCall;
@@ -1849,4 +1862,199 @@ public void testInitHttpClientWithCustomHttpClient() throws Exception {
18491862
client.httpClient().interceptors().stream()
18501863
.anyMatch(interceptor -> interceptor instanceof RetryInterceptor));
18511864
}
1865+
1866+
/** Credentials that issue whatever token the test sets, refreshing on every use. */
1867+
private static final class SwitchableCredentials extends GoogleCredentials {
1868+
final AtomicReference<String> token;
1869+
volatile boolean failRefresh = false;
1870+
1871+
SwitchableCredentials(String token) {
1872+
this.token = new AtomicReference<>(token);
1873+
}
1874+
1875+
@Override
1876+
public AccessToken refreshAccessToken() throws IOException {
1877+
if (failRefresh) {
1878+
throw new IOException("Token endpoint unavailable.");
1879+
}
1880+
// Expiring inside the refresh margin makes every refreshIfExpired() call refresh again.
1881+
return AccessToken.newBuilder()
1882+
.setTokenValue(token.get())
1883+
.setExpirationTime(new Date(System.currentTimeMillis() + 60_000))
1884+
.build();
1885+
}
1886+
}
1887+
1888+
private static HttpApiClient vertexClient(
1889+
GoogleCredentials credentials, Optional<ClientOptions> clientOptions) {
1890+
return new HttpApiClient(
1891+
Optional.empty(),
1892+
Optional.of(PROJECT),
1893+
Optional.of(LOCATION),
1894+
Optional.of(credentials),
1895+
Optional.of(FAST_RETRY_HTTP_OPTIONS),
1896+
clientOptions);
1897+
}
1898+
1899+
/** Appends a fake server after the SDK interceptors, so each attempt is answered locally. */
1900+
private static void addFakeServer(HttpApiClient client, Interceptor fakeServer)
1901+
throws Exception {
1902+
Field clientField = ApiClient.class.getDeclaredField("httpClient");
1903+
clientField.setAccessible(true);
1904+
clientField.set(client, client.httpClient().newBuilder().addInterceptor(fakeServer).build());
1905+
}
1906+
1907+
private static Response respond(Request request, int code) {
1908+
return new Response.Builder()
1909+
.request(request)
1910+
.protocol(Protocol.HTTP_1_1)
1911+
.code(code)
1912+
.message("")
1913+
.body(ResponseBody.create("{}", null))
1914+
.build();
1915+
}
1916+
1917+
@Test
1918+
public void testRequest_retryAfterTokenRotates_sendsCurrentToken() throws Exception {
1919+
SwitchableCredentials credentials = new SwitchableCredentials("token-a");
1920+
HttpApiClient client = vertexClient(credentials, Optional.empty());
1921+
List<Request> attempts = new ArrayList<>();
1922+
addFakeServer(
1923+
client,
1924+
chain -> {
1925+
attempts.add(chain.request());
1926+
if (attempts.size() == 1) {
1927+
credentials.token.set("token-b");
1928+
return respond(chain.request(), 503);
1929+
}
1930+
return respond(chain.request(), 200);
1931+
});
1932+
1933+
client.request("POST", TEST_PATH, TEST_REQUEST_JSON, Optional.empty()).close();
1934+
1935+
assertEquals(2, attempts.size());
1936+
assertEquals("Bearer token-a", attempts.get(0).header("Authorization"));
1937+
assertEquals("Bearer token-b", attempts.get(1).header("Authorization"));
1938+
}
1939+
1940+
@Test
1941+
public void testAsyncRequest_queuedBehindDispatcherLimit_sendsTokenCurrentAtSendTime()
1942+
throws Exception {
1943+
SwitchableCredentials credentials = new SwitchableCredentials("token-a");
1944+
HttpApiClient client =
1945+
vertexClient(
1946+
credentials, Optional.of(ClientOptions.builder().maxConnectionsPerHost(1).build()));
1947+
CountDownLatch firstCallStarted = new CountDownLatch(1);
1948+
CountDownLatch releaseFirstCall = new CountDownLatch(1);
1949+
List<Request> attempts = Collections.synchronizedList(new ArrayList<>());
1950+
addFakeServer(
1951+
client,
1952+
chain -> {
1953+
attempts.add(chain.request());
1954+
if (attempts.size() == 1) {
1955+
firstCallStarted.countDown();
1956+
try {
1957+
releaseFirstCall.await();
1958+
} catch (InterruptedException e) {
1959+
throw new IOException(e);
1960+
}
1961+
}
1962+
return respond(chain.request(), 200);
1963+
});
1964+
1965+
CompletableFuture<ApiResponse> first =
1966+
client.asyncRequest("POST", TEST_PATH, TEST_REQUEST_JSON, Optional.empty());
1967+
assertTrue(firstCallStarted.await(5, TimeUnit.SECONDS));
1968+
// Built with token-a, then queued until the first call finishes.
1969+
CompletableFuture<ApiResponse> second =
1970+
client.asyncRequest("POST", TEST_PATH, TEST_REQUEST_JSON, Optional.empty());
1971+
credentials.token.set("token-b");
1972+
releaseFirstCall.countDown();
1973+
first.get(5, TimeUnit.SECONDS).close();
1974+
second.get(5, TimeUnit.SECONDS).close();
1975+
1976+
assertEquals(2, attempts.size());
1977+
assertEquals("Bearer token-b", attempts.get(1).header("Authorization"));
1978+
}
1979+
1980+
@Test
1981+
public void testRequest_authorizationChangedByCustomClient_isKeptOnRetry() throws Exception {
1982+
SwitchableCredentials credentials = new SwitchableCredentials("token-a");
1983+
OkHttpClient customClient =
1984+
new OkHttpClient.Builder()
1985+
.addInterceptor(
1986+
chain ->
1987+
chain.proceed(
1988+
chain
1989+
.request()
1990+
.newBuilder()
1991+
.header("Authorization", "Bearer caller-token")
1992+
.build()))
1993+
.build();
1994+
HttpApiClient client =
1995+
vertexClient(
1996+
credentials, Optional.of(ClientOptions.builder().customHttpClient(customClient).build()));
1997+
List<Request> attempts = new ArrayList<>();
1998+
addFakeServer(
1999+
client,
2000+
chain -> {
2001+
attempts.add(chain.request());
2002+
if (attempts.size() == 1) {
2003+
credentials.token.set("token-b");
2004+
return respond(chain.request(), 503);
2005+
}
2006+
return respond(chain.request(), 200);
2007+
});
2008+
2009+
client.request("POST", TEST_PATH, TEST_REQUEST_JSON, Optional.empty()).close();
2010+
2011+
assertEquals(2, attempts.size());
2012+
assertEquals("Bearer caller-token", attempts.get(0).header("Authorization"));
2013+
assertEquals("Bearer caller-token", attempts.get(1).header("Authorization"));
2014+
}
2015+
2016+
@Test
2017+
public void testRequest_apiKeyRetry_sendsNoAuthorizationHeader() throws Exception {
2018+
HttpApiClient client =
2019+
new HttpApiClient(
2020+
Optional.of(API_KEY), Optional.of(FAST_RETRY_HTTP_OPTIONS), Optional.empty());
2021+
List<Request> attempts = new ArrayList<>();
2022+
addFakeServer(
2023+
client,
2024+
chain -> {
2025+
attempts.add(chain.request());
2026+
return respond(chain.request(), attempts.size() == 1 ? 503 : 200);
2027+
});
2028+
2029+
client.request("POST", TEST_PATH, TEST_REQUEST_JSON, Optional.empty()).close();
2030+
2031+
assertEquals(2, attempts.size());
2032+
for (Request attempt : attempts) {
2033+
assertEquals(API_KEY, attempt.header("x-goog-api-key"));
2034+
assertNull(attempt.header("Authorization"));
2035+
}
2036+
}
2037+
2038+
@Test
2039+
public void testAsyncRequest_refreshFailsOnRetry_completesWithRefreshError() throws Exception {
2040+
SwitchableCredentials credentials = new SwitchableCredentials("token-a");
2041+
HttpApiClient client = vertexClient(credentials, Optional.empty());
2042+
addFakeServer(
2043+
client,
2044+
chain -> {
2045+
credentials.failRefresh = true;
2046+
return respond(chain.request(), 503);
2047+
});
2048+
2049+
ExecutionException exception =
2050+
assertThrows(
2051+
ExecutionException.class,
2052+
() ->
2053+
client
2054+
.asyncRequest("POST", TEST_PATH, TEST_REQUEST_JSON, Optional.empty())
2055+
.get(5, TimeUnit.SECONDS));
2056+
2057+
assertTrue(exception.getCause() instanceof GenAiIOException);
2058+
assertEquals("Failed to refresh credentials.", exception.getCause().getCause().getMessage());
2059+
}
18522060
}

0 commit comments

Comments
 (0)