diff --git a/.castiron.stats.yml b/.castiron.stats.yml index 2a9336c1d..c0fce825b 100644 --- a/.castiron.stats.yml +++ b/.castiron.stats.yml @@ -5,4 +5,4 @@ openapi_transformed_spec_hash: 76df9e32111a39f4e8916e2b40ccb12a config_hash: 0adc48b9b2b37cb31efd58e02e66aaa8 codegen_sha: e26b222ae4e5d3ded59c470aff3a5e3ce5be6e82 codegen_hash: 15eb661132463d5a980394613a4b5652ad12ae2ac500da1f725af3c351bcbb38 -public_codegen_sha: d04db21ff39146f727b27bdedc013beefb9b7b10 +public_codegen_sha: 41afd676aeb6db766d2d39622aef60c8e70c34ca diff --git a/openai-java-core/src/main/kotlin/com/openai/auth/WorkloadIdentityAuth.kt b/openai-java-core/src/main/kotlin/com/openai/auth/WorkloadIdentityAuth.kt index c22ce5ad2..e23b4eb75 100644 --- a/openai-java-core/src/main/kotlin/com/openai/auth/WorkloadIdentityAuth.kt +++ b/openai-java-core/src/main/kotlin/com/openai/auth/WorkloadIdentityAuth.kt @@ -155,7 +155,7 @@ internal class WorkloadIdentityAuth( return when (action) { is TokenAction.ReturnCached -> CompletableFuture.completedFuture(action.token) is TokenAction.BackgroundRefresh -> { - performRefreshAndComplete(action.future) + startAsyncRefresh(action.future) CompletableFuture.completedFuture(action.token) } is TokenAction.WaitForRefresh -> @@ -165,20 +165,25 @@ internal class WorkloadIdentityAuth( is TokenRefreshResult.Failure -> throw result.error } } - is TokenAction.ForegroundRefresh -> { - val refresh = refreshTokenAsync() - refresh.whenComplete { token, error -> - finishRefresh(action.future, token, unwrapCompletionException(error)) - } - refresh - } + is TokenAction.ForegroundRefresh -> startAsyncRefresh(action.future) } } - private fun performRefreshAndComplete(future: CompletableFuture) { - refreshTokenAsync().whenComplete { token, error -> + private fun startAsyncRefresh( + future: CompletableFuture + ): CompletableFuture { + val refresh = + try { + refreshTokenAsync() + } catch (error: Throwable) { + finishRefresh(future, null, error) + return CompletableFuture().also { it.completeExceptionally(error) } + } + + refresh.whenComplete { token, error -> finishRefresh(future, token, unwrapCompletionException(error)) } + return refresh } private fun finishRefresh( diff --git a/openai-java-core/src/test/kotlin/com/openai/auth/WorkloadIdentityAuthSynchronousProviderFailureTest.kt b/openai-java-core/src/test/kotlin/com/openai/auth/WorkloadIdentityAuthSynchronousProviderFailureTest.kt new file mode 100644 index 000000000..ad9a51eca --- /dev/null +++ b/openai-java-core/src/test/kotlin/com/openai/auth/WorkloadIdentityAuthSynchronousProviderFailureTest.kt @@ -0,0 +1,199 @@ +package com.openai.auth + +import com.fasterxml.jackson.databind.json.JsonMapper +import com.openai.core.http.HttpClient +import com.openai.core.http.HttpRequest +import com.openai.core.http.HttpResponse +import java.io.ByteArrayInputStream +import java.util.concurrent.CompletableFuture +import java.util.concurrent.CountDownLatch +import java.util.concurrent.ExecutionException +import java.util.concurrent.Executors +import java.util.concurrent.TimeUnit.SECONDS +import org.assertj.core.api.Assertions.assertThat +import org.junit.jupiter.api.Test +import org.junit.jupiter.api.Timeout +import org.junit.jupiter.api.assertThrows +import org.junit.jupiter.params.ParameterizedTest +import org.junit.jupiter.params.provider.ValueSource +import org.mockito.kotlin.any +import org.mockito.kotlin.mock +import org.mockito.kotlin.times +import org.mockito.kotlin.verify +import org.mockito.kotlin.verifyNoInteractions +import org.mockito.kotlin.whenever + +@Timeout(15) +internal class WorkloadIdentityAuthSynchronousProviderFailureTest { + + @Test + fun getTokenAsync_clearsForegroundRefreshAfterSynchronousProviderFailure() { + val failure = IllegalStateException("provider failed") + val provider = mock() + whenever(provider.getTokenAsync(any(), any())).thenThrow(failure) + val httpClient = mock() + val auth = createAuth(provider, httpClient) + + assertFailure(auth.getTokenAsync(), failure) + assertFailure(auth.getTokenAsync(), failure) + + verify(provider, times(2)).getTokenAsync(any(), any()) + verifyNoInteractions(httpClient) + } + + @Test + fun getTokenAsync_keepsCachedTokenAndAllowsAnotherBackgroundRefreshAfterSynchronousFailure() { + val provider = mock() + whenever(provider.tokenType()).thenReturn(SubjectTokenType.JWT) + whenever(provider.getToken(any(), any())).thenReturn("subject-token") + whenever(provider.getTokenAsync(any(), any())) + .thenThrow(IllegalStateException("provider failed")) + val httpClient = mock() + val cachedResponse = mockTokenResponse("cached-token", expiresIn = 600) + whenever(httpClient.execute(any())).thenReturn(cachedResponse) + val auth = createAuth(provider, httpClient) + + assertThat(auth.getToken()).isEqualTo("cached-token") + assertThat(auth.getTokenAsync().get(5, SECONDS)).isEqualTo("cached-token") + assertThat(auth.getTokenAsync().get(5, SECONDS)).isEqualTo("cached-token") + + verify(provider, times(2)).getTokenAsync(any(), any()) + verify(httpClient, times(1)).execute(any()) + } + + @ParameterizedTest + @ValueSource(booleans = [true, false]) + fun getTokenAsync_recoversAfterProviderFailure(synchronousFailure: Boolean) { + val failure = IllegalStateException("provider failed") + val provider = mock() + whenever(provider.tokenType()).thenReturn(SubjectTokenType.JWT) + whenever(provider.getTokenAsync(any(), any())) + .thenAnswer { + if (synchronousFailure) throw failure + CompletableFuture().also { it.completeExceptionally(failure) } + } + .thenReturn(CompletableFuture.completedFuture("subject-token")) + val httpClient = mock() + val recoveredResponse = mockTokenResponse("recovered-token") + whenever(httpClient.executeAsync(any())) + .thenReturn(CompletableFuture.completedFuture(recoveredResponse)) + val auth = createAuth(provider, httpClient) + + assertFailure(auth.getTokenAsync(), failure) + assertThat(auth.getTokenAsync().get(5, SECONDS)).isEqualTo("recovered-token") + assertThat(auth.getTokenAsync().get(5, SECONDS)).isEqualTo("recovered-token") + + verify(provider, times(2)).getTokenAsync(any(), any()) + verify(httpClient, times(1)).executeAsync(any()) + } + + @Test + fun getTokenAsync_recoversBackgroundRefreshWithoutDiscardingCachedToken() { + val provider = mock() + whenever(provider.tokenType()).thenReturn(SubjectTokenType.JWT) + whenever(provider.getToken(any(), any())).thenReturn("subject-token") + whenever(provider.getTokenAsync(any(), any())) + .thenThrow(IllegalStateException("provider failed")) + .thenReturn(CompletableFuture.completedFuture("subject-token")) + val httpClient = mock() + val cachedResponse = mockTokenResponse("cached-token", expiresIn = 600) + whenever(httpClient.execute(any())).thenReturn(cachedResponse) + val exchange = CompletableFuture() + whenever(httpClient.executeAsync(any())).thenReturn(exchange) + val auth = createAuth(provider, httpClient) + + assertThat(auth.getToken()).isEqualTo("cached-token") + // A failed refresh leaves the valid token available. + assertThat(auth.getTokenAsync().get(5, SECONDS)).isEqualTo("cached-token") + // The next call starts recovery; concurrent requests keep using the cached token. + assertThat(auth.getTokenAsync().get(5, SECONDS)).isEqualTo("cached-token") + assertThat(auth.getTokenAsync().get(5, SECONDS)).isEqualTo("cached-token") + + exchange.complete(mockTokenResponse("recovered-token")) + assertThat(auth.getTokenAsync().get(5, SECONDS)).isEqualTo("recovered-token") + + verify(provider, times(2)).getTokenAsync(any(), any()) + verify(httpClient, times(1)).executeAsync(any()) + } + + @Test + fun getTokenAsync_settlesConcurrentWaitersAfterSynchronousProviderFailure() { + val failure = IllegalStateException("provider failed") + val providerStarted = CountDownLatch(1) + val releaseProvider = CountDownLatch(1) + val provider = mock() + whenever(provider.tokenType()).thenReturn(SubjectTokenType.JWT) + whenever(provider.getTokenAsync(any(), any())) + .thenAnswer { + providerStarted.countDown() + check(releaseProvider.await(5, SECONDS)) { "Provider was not released" } + throw failure + } + .thenReturn(CompletableFuture.completedFuture("subject-token")) + val httpClient = mock() + val recoveredResponse = mockTokenResponse("recovered-token") + whenever(httpClient.executeAsync(any())) + .thenReturn(CompletableFuture.completedFuture(recoveredResponse)) + val auth = createAuth(provider, httpClient) + val executor = Executors.newSingleThreadExecutor() + + try { + val initiatingCall = executor.submit> { auth.getTokenAsync() } + assertThat(providerStarted.await(5, SECONDS)).isTrue() + val waiters = List(3) { auth.getTokenAsync() } + waiters.forEach { assertThat(it).isNotDone() } + verify(provider, times(1)).getTokenAsync(any(), any()) + + releaseProvider.countDown() + assertFailure(initiatingCall.get(5, SECONDS), failure) + waiters.forEach { assertFailure(it, failure) } + + assertThat(auth.getTokenAsync().get(5, SECONDS)).isEqualTo("recovered-token") + verify(provider, times(2)).getTokenAsync(any(), any()) + verify(httpClient, times(1)).executeAsync(any()) + } finally { + releaseProvider.countDown() + executor.shutdownNow() + assertThat(executor.awaitTermination(5, SECONDS)).isTrue() + } + } + + private fun assertFailure(future: CompletableFuture, failure: Throwable) { + val thrown = assertThrows { future.get(5, SECONDS) } + assertThat(thrown.cause).isSameAs(failure) + } + + private fun createAuth( + provider: SubjectTokenProvider, + httpClient: HttpClient, + ): WorkloadIdentityAuth = + WorkloadIdentityAuth( + config = + WorkloadIdentity.builder() + .clientId("client-id") + .identityProviderId("provider-id") + .serviceAccountId("service-account-id") + .provider(provider) + .refreshBufferSeconds(1200) + .build(), + httpClient = httpClient, + jsonMapper = JsonMapper(), + ) + + private fun mockTokenResponse(token: String, expiresIn: Int = 3600): HttpResponse { + val body = + """ + { + "access_token": "$token", + "issued_token_type": "urn:ietf:params:oauth:token-type:access_token", + "token_type": "Bearer", + "expires_in": $expiresIn + } + """ + .trimIndent() + val response = mock() + whenever(response.statusCode()).thenReturn(200) + whenever(response.body()).thenAnswer { ByteArrayInputStream(body.toByteArray()) } + return response + } +}