From 220c4ded7534eb679b4e0fa37425f454cae88d90 Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Tue, 25 Aug 2026 05:40:57 +0000 Subject: [PATCH 1/5] feat: integrate x509 clients --- .../com/openai/client/okhttp/OkHttpClient.kt | 67 ++- .../client/okhttp/OpenAIOkHttpClient.kt | 88 +++- .../client/okhttp/OpenAIOkHttpClientAsync.kt | 88 +++- .../client/okhttp/X509ClientIntegration.kt | 136 ++++++ .../com/openai/client/okhttp/X509Transport.kt | 12 +- .../openai/client/okhttp/OkHttpClientTest.kt | 55 +++ .../okhttp/OpenAIOkHttpClientX509Test.kt | 415 ++++++++++++++++++ .../okhttp/X509WireTestInfrastructure.kt | 7 +- .../kotlin/com/openai/core/ClientOptions.kt | 122 ++++- .../com/openai/core/RequestAuthentication.kt | 59 +++ .../com/openai/core/ClientOptionsTest.kt | 67 +++ 11 files changed, 1048 insertions(+), 68 deletions(-) create mode 100644 openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/X509ClientIntegration.kt create mode 100644 openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/OpenAIOkHttpClientX509Test.kt create mode 100644 openai-java-core/src/main/kotlin/com/openai/core/RequestAuthentication.kt diff --git a/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/OkHttpClient.kt b/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/OkHttpClient.kt index 367b1ead4..160705a41 100644 --- a/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/OkHttpClient.kt +++ b/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/OkHttpClient.kt @@ -17,8 +17,10 @@ import java.net.Proxy import java.time.Duration import java.util.concurrent.CancellationException import java.util.concurrent.CompletableFuture +import java.util.concurrent.ConcurrentHashMap import java.util.concurrent.ExecutorService import java.util.concurrent.TimeUnit +import java.util.concurrent.atomic.AtomicBoolean import javax.net.ssl.HostnameVerifier import javax.net.ssl.SSLSocketFactory import javax.net.ssl.X509TrustManager @@ -27,6 +29,7 @@ import okhttp3.Call import okhttp3.Callback import okhttp3.ConnectionPool import okhttp3.Dispatcher +import okhttp3.EventListener import okhttp3.HttpUrl import okhttp3.HttpUrl.Companion.toHttpUrl import okhttp3.MediaType @@ -40,11 +43,13 @@ import okio.buffer import okio.sink class OkHttpClient -internal constructor(@JvmSynthetic internal val okHttpClient: okhttp3.OkHttpClient) : HttpClient { +private constructor( + @get:JvmSynthetic internal val okHttpClient: okhttp3.OkHttpClient, + private val callTracker: CallTracker, +) : HttpClient { override fun execute(request: HttpRequest, requestOptions: RequestOptions): HttpResponse { val call = newCall(request, requestOptions) - return try { call.execute().toHttpResponse() } catch (e: IOException) { @@ -73,7 +78,7 @@ internal constructor(@JvmSynthetic internal val okHttpClient: okhttp3.OkHttpClie } ) - future.whenComplete { _, e -> + future.whenComplete { response, e -> if (e is CancellationException) { call.cancel() } @@ -84,12 +89,15 @@ internal constructor(@JvmSynthetic internal val okHttpClient: okhttp3.OkHttpClie } override fun close() { - okHttpClient.dispatcher.executorService.shutdown() - okHttpClient.connectionPool.evictAll() - okHttpClient.cache?.close() + if (callTracker.close()) { + okHttpClient.dispatcher.executorService.shutdown() + okHttpClient.connectionPool.evictAll() + okHttpClient.cache?.close() + } } private fun newCall(request: HttpRequest, requestOptions: RequestOptions): Call { + callTracker.ensureOpen() val clientBuilder = okHttpClient.newBuilder() requestOptions.timeout?.let { @@ -174,11 +182,13 @@ internal constructor(@JvmSynthetic internal val okHttpClient: okhttp3.OkHttpClie this.hostnameVerifier = hostnameVerifier } - fun build(): OkHttpClient = - OkHttpClient( + fun build(): OkHttpClient { + val callTracker = CallTracker() + return OkHttpClient( okhttp3.OkHttpClient.Builder() // `RetryingHttpClient` handles retries if the user enabled them. .retryOnConnectionFailure(false) + .eventListenerFactory(callTracker.eventListenerFactory) .followRedirects(followRedirects) .followSslRedirects(followRedirects) .connectTimeout(timeout.connect()) @@ -235,8 +245,47 @@ internal constructor(@JvmSynthetic internal val okHttpClient: okhttp3.OkHttpClie // We usually make all our requests to the same host so it makes sense to // raise the per-host limit to the overall limit. dispatcher.maxRequestsPerHost = dispatcher.maxRequests - } + }, + callTracker, ) + } + } +} + +private class CallTracker { + private val closed = AtomicBoolean() + private val activeCalls = ConcurrentHashMap.newKeySet() + + val eventListenerFactory = + EventListener.Factory { + object : EventListener() { + override fun callStart(call: Call) { + activeCalls.add(call) + if (closed.get()) { + call.cancel() + } + } + + override fun callEnd(call: Call) { + activeCalls.remove(call) + } + + override fun callFailed(call: Call, ioe: IOException) { + activeCalls.remove(call) + } + } + } + + fun ensureOpen() { + check(!closed.get()) { "HTTP client is closed" } + } + + fun close(): Boolean { + if (!closed.compareAndSet(false, true)) { + return false + } + activeCalls.forEach(Call::cancel) + return true } } diff --git a/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/OpenAIOkHttpClient.kt b/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/OpenAIOkHttpClient.kt index 6539ce5dd..c5b363994 100644 --- a/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/OpenAIOkHttpClient.kt +++ b/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/OpenAIOkHttpClient.kt @@ -2,6 +2,7 @@ package com.openai.client.okhttp import com.fasterxml.jackson.databind.json.JsonMapper import com.openai.auth.WorkloadIdentity +import com.openai.auth.X509WorkloadIdentity import com.openai.azure.AzureOpenAIServiceVersion import com.openai.azure.AzureUrlPathMode import com.openai.client.OpenAIClient @@ -40,6 +41,30 @@ class OpenAIOkHttpClient private constructor() { /** Returns a mutable builder for constructing an instance of [OpenAIClient]. */ @JvmStatic fun builder() = Builder() + /** + * Preview: returns a builder for direct-only X.509 workload identity federation. + * + * The token issuer and API base URL are fixed OpenAI mTLS endpoints. Each built client owns + * a newly bound session from [transport] and presents its fixed certificate alias on both + * network legs. + */ + @JvmStatic + fun x509Builder(identity: X509WorkloadIdentity, transport: X509Transport) = + Builder.x509(X509ClientConfiguration.create(identity, transport::bind)) + + @JvmSynthetic + internal fun x509BuilderForTest( + identity: X509WorkloadIdentity, + transport: X509Transport, + exchangeProxy: Proxy, + apiProxy: Proxy, + ) = + Builder.x509( + X509ClientConfiguration.create(identity) { timeout -> + transport.bindForTest(timeout, exchangeProxy, apiProxy) + } + ) + /** * Returns a client configured using system properties and environment variables. * @@ -51,7 +76,13 @@ class OpenAIOkHttpClient private constructor() { /** A builder for [OpenAIOkHttpClient]. */ class Builder internal constructor() { + companion object { + @JvmSynthetic + internal fun x509(configuration: X509ClientConfiguration) = Builder(configuration) + } + private var clientOptions: ClientOptions.Builder = ClientOptions.builder() + private var x509Configuration: X509ClientConfiguration? = null private var dispatcherExecutorService: ExecutorService? = null private var followRedirects: Boolean = true private var proxy: Proxy? = null @@ -62,6 +93,11 @@ class OpenAIOkHttpClient private constructor() { private var trustManager: X509TrustManager? = null private var hostnameVerifier: HostnameVerifier? = null + private constructor(x509Configuration: X509ClientConfiguration) : this() { + this.x509Configuration = x509Configuration + x509Configuration.reserve(clientOptions) + } + /** * The executor service to use for running HTTP requests. * @@ -71,6 +107,7 @@ class OpenAIOkHttpClient private constructor() { * This class takes ownership of the executor service and shuts it down when closed. */ fun dispatcherExecutorService(dispatcherExecutorService: ExecutorService?) = apply { + requireGenericTransport("dispatcherExecutorService") this.dispatcherExecutorService = dispatcherExecutorService } @@ -80,6 +117,7 @@ class OpenAIOkHttpClient private constructor() { * Defaults to true. */ fun followRedirects(followRedirects: Boolean) = apply { + requireGenericTransport("followRedirects") this.followRedirects = followRedirects } @@ -90,7 +128,10 @@ class OpenAIOkHttpClient private constructor() { fun dispatcherExecutorService(dispatcherExecutorService: Optional) = dispatcherExecutorService(dispatcherExecutorService.getOrNull()) - fun proxy(proxy: Proxy?) = apply { this.proxy = proxy } + fun proxy(proxy: Proxy?) = apply { + requireGenericTransport("proxy") + this.proxy = proxy + } /** Alias for calling [Builder.proxy] with `proxy.orElse(null)`. */ fun proxy(proxy: Optional) = proxy(proxy.getOrNull()) @@ -100,6 +141,7 @@ class OpenAIOkHttpClient private constructor() { * Required`. */ fun proxyAuthenticator(proxyAuthenticator: ProxyAuthenticator?) = apply { + requireGenericTransport("proxyAuthenticator") this.proxyAuthenticator = proxyAuthenticator } @@ -117,6 +159,7 @@ class OpenAIOkHttpClient private constructor() { * If unset, then OkHttp's default is used. */ fun maxIdleConnections(maxIdleConnections: Int?) = apply { + requireGenericTransport("maxIdleConnections") this.maxIdleConnections = maxIdleConnections } @@ -142,6 +185,7 @@ class OpenAIOkHttpClient private constructor() { * If unset, then OkHttp's default is used. */ fun keepAliveDuration(keepAliveDuration: Duration?) = apply { + requireGenericTransport("keepAliveDuration") this.keepAliveDuration = keepAliveDuration } @@ -159,6 +203,7 @@ class OpenAIOkHttpClient private constructor() { * lost if the implementation is modified. */ fun sslSocketFactory(sslSocketFactory: SSLSocketFactory?) = apply { + requireGenericTransport("sslSocketFactory") this.sslSocketFactory = sslSocketFactory } @@ -176,6 +221,7 @@ class OpenAIOkHttpClient private constructor() { * lost if the implementation is modified. */ fun trustManager(trustManager: X509TrustManager?) = apply { + requireGenericTransport("trustManager") this.trustManager = trustManager } @@ -190,6 +236,7 @@ class OpenAIOkHttpClient private constructor() { * If unset, then a default hostname verifier is used. */ fun hostnameVerifier(hostnameVerifier: HostnameVerifier?) = apply { + requireGenericTransport("hostnameVerifier") this.hostnameVerifier = hostnameVerifier } @@ -472,22 +519,29 @@ class OpenAIOkHttpClient private constructor() { */ fun build(): OpenAIClient = OpenAIClientImpl( - clientOptions - .httpClient( - OkHttpClient.builder() - .timeout(clientOptions.timeout()) - .followRedirects(followRedirects) - .proxy(proxy) - .proxyAuthenticator(proxyAuthenticator) - .maxIdleConnections(maxIdleConnections) - .keepAliveDuration(keepAliveDuration) - .dispatcherExecutorService(dispatcherExecutorService) - .sslSocketFactory(sslSocketFactory) - .trustManager(trustManager) - .hostnameVerifier(hostnameVerifier) - .build() - ) - .build() + x509Configuration?.buildClientOptions(clientOptions) + ?: clientOptions + .httpClient( + OkHttpClient.builder() + .timeout(clientOptions.timeout()) + .followRedirects(followRedirects) + .proxy(proxy) + .proxyAuthenticator(proxyAuthenticator) + .maxIdleConnections(maxIdleConnections) + .keepAliveDuration(keepAliveDuration) + .dispatcherExecutorService(dispatcherExecutorService) + .sslSocketFactory(sslSocketFactory) + .trustManager(trustManager) + .hostnameVerifier(hostnameVerifier) + .build() + ) + .build() ) + + private fun requireGenericTransport(option: String) { + require(x509Configuration == null) { + "$option cannot be configured on an X.509 client builder" + } + } } } diff --git a/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/OpenAIOkHttpClientAsync.kt b/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/OpenAIOkHttpClientAsync.kt index af1afa25f..9983e66d0 100644 --- a/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/OpenAIOkHttpClientAsync.kt +++ b/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/OpenAIOkHttpClientAsync.kt @@ -2,6 +2,7 @@ package com.openai.client.okhttp import com.fasterxml.jackson.databind.json.JsonMapper import com.openai.auth.WorkloadIdentity +import com.openai.auth.X509WorkloadIdentity import com.openai.azure.AzureOpenAIServiceVersion import com.openai.azure.AzureUrlPathMode import com.openai.client.OpenAIClientAsync @@ -40,6 +41,30 @@ class OpenAIOkHttpClientAsync private constructor() { /** Returns a mutable builder for constructing an instance of [OpenAIClientAsync]. */ @JvmStatic fun builder() = Builder() + /** + * Preview: returns a builder for direct-only X.509 workload identity federation. + * + * The token issuer and API base URL are fixed OpenAI mTLS endpoints. Each built client owns + * a newly bound session from [transport] and presents its fixed certificate alias on both + * network legs. + */ + @JvmStatic + fun x509Builder(identity: X509WorkloadIdentity, transport: X509Transport) = + Builder.x509(X509ClientConfiguration.create(identity, transport::bind)) + + @JvmSynthetic + internal fun x509BuilderForTest( + identity: X509WorkloadIdentity, + transport: X509Transport, + exchangeProxy: Proxy, + apiProxy: Proxy, + ) = + Builder.x509( + X509ClientConfiguration.create(identity) { timeout -> + transport.bindForTest(timeout, exchangeProxy, apiProxy) + } + ) + /** * Returns a client configured using system properties and environment variables. * @@ -51,7 +76,13 @@ class OpenAIOkHttpClientAsync private constructor() { /** A builder for [OpenAIOkHttpClientAsync]. */ class Builder internal constructor() { + companion object { + @JvmSynthetic + internal fun x509(configuration: X509ClientConfiguration) = Builder(configuration) + } + private var clientOptions: ClientOptions.Builder = ClientOptions.builder() + private var x509Configuration: X509ClientConfiguration? = null private var dispatcherExecutorService: ExecutorService? = null private var followRedirects: Boolean = true private var proxy: Proxy? = null @@ -62,6 +93,11 @@ class OpenAIOkHttpClientAsync private constructor() { private var trustManager: X509TrustManager? = null private var hostnameVerifier: HostnameVerifier? = null + private constructor(x509Configuration: X509ClientConfiguration) : this() { + this.x509Configuration = x509Configuration + x509Configuration.reserve(clientOptions) + } + /** * The executor service to use for running HTTP requests. * @@ -71,6 +107,7 @@ class OpenAIOkHttpClientAsync private constructor() { * This class takes ownership of the executor service and shuts it down when closed. */ fun dispatcherExecutorService(dispatcherExecutorService: ExecutorService?) = apply { + requireGenericTransport("dispatcherExecutorService") this.dispatcherExecutorService = dispatcherExecutorService } @@ -80,6 +117,7 @@ class OpenAIOkHttpClientAsync private constructor() { * Defaults to true. */ fun followRedirects(followRedirects: Boolean) = apply { + requireGenericTransport("followRedirects") this.followRedirects = followRedirects } @@ -90,7 +128,10 @@ class OpenAIOkHttpClientAsync private constructor() { fun dispatcherExecutorService(dispatcherExecutorService: Optional) = dispatcherExecutorService(dispatcherExecutorService.getOrNull()) - fun proxy(proxy: Proxy?) = apply { this.proxy = proxy } + fun proxy(proxy: Proxy?) = apply { + requireGenericTransport("proxy") + this.proxy = proxy + } /** Alias for calling [Builder.proxy] with `proxy.orElse(null)`. */ fun proxy(proxy: Optional) = proxy(proxy.getOrNull()) @@ -100,6 +141,7 @@ class OpenAIOkHttpClientAsync private constructor() { * Required`. */ fun proxyAuthenticator(proxyAuthenticator: ProxyAuthenticator?) = apply { + requireGenericTransport("proxyAuthenticator") this.proxyAuthenticator = proxyAuthenticator } @@ -117,6 +159,7 @@ class OpenAIOkHttpClientAsync private constructor() { * If unset, then OkHttp's default is used. */ fun maxIdleConnections(maxIdleConnections: Int?) = apply { + requireGenericTransport("maxIdleConnections") this.maxIdleConnections = maxIdleConnections } @@ -142,6 +185,7 @@ class OpenAIOkHttpClientAsync private constructor() { * If unset, then OkHttp's default is used. */ fun keepAliveDuration(keepAliveDuration: Duration?) = apply { + requireGenericTransport("keepAliveDuration") this.keepAliveDuration = keepAliveDuration } @@ -159,6 +203,7 @@ class OpenAIOkHttpClientAsync private constructor() { * lost if the implementation is modified. */ fun sslSocketFactory(sslSocketFactory: SSLSocketFactory?) = apply { + requireGenericTransport("sslSocketFactory") this.sslSocketFactory = sslSocketFactory } @@ -176,6 +221,7 @@ class OpenAIOkHttpClientAsync private constructor() { * lost if the implementation is modified. */ fun trustManager(trustManager: X509TrustManager?) = apply { + requireGenericTransport("trustManager") this.trustManager = trustManager } @@ -190,6 +236,7 @@ class OpenAIOkHttpClientAsync private constructor() { * If unset, then a default hostname verifier is used. */ fun hostnameVerifier(hostnameVerifier: HostnameVerifier?) = apply { + requireGenericTransport("hostnameVerifier") this.hostnameVerifier = hostnameVerifier } @@ -472,22 +519,29 @@ class OpenAIOkHttpClientAsync private constructor() { */ fun build(): OpenAIClientAsync = OpenAIClientAsyncImpl( - clientOptions - .httpClient( - OkHttpClient.builder() - .timeout(clientOptions.timeout()) - .followRedirects(followRedirects) - .proxy(proxy) - .proxyAuthenticator(proxyAuthenticator) - .maxIdleConnections(maxIdleConnections) - .keepAliveDuration(keepAliveDuration) - .dispatcherExecutorService(dispatcherExecutorService) - .sslSocketFactory(sslSocketFactory) - .trustManager(trustManager) - .hostnameVerifier(hostnameVerifier) - .build() - ) - .build() + x509Configuration?.buildClientOptions(clientOptions) + ?: clientOptions + .httpClient( + OkHttpClient.builder() + .timeout(clientOptions.timeout()) + .followRedirects(followRedirects) + .proxy(proxy) + .proxyAuthenticator(proxyAuthenticator) + .maxIdleConnections(maxIdleConnections) + .keepAliveDuration(keepAliveDuration) + .dispatcherExecutorService(dispatcherExecutorService) + .sslSocketFactory(sslSocketFactory) + .trustManager(trustManager) + .hostnameVerifier(hostnameVerifier) + .build() + ) + .build() ) + + private fun requireGenericTransport(option: String) { + require(x509Configuration == null) { + "$option cannot be configured on an X.509 client builder" + } + } } } diff --git a/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/X509ClientIntegration.kt b/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/X509ClientIntegration.kt new file mode 100644 index 000000000..0bf55da42 --- /dev/null +++ b/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/X509ClientIntegration.kt @@ -0,0 +1,136 @@ +package com.openai.client.okhttp + +import com.openai.auth.X509WorkloadIdentity +import com.openai.core.ClientOptions +import com.openai.core.Timeout +import com.openai.core.http.HttpRequest +import com.openai.core.http.HttpRequestAuthenticator +import java.util.Locale +import java.util.concurrent.CancellationException +import java.util.concurrent.CompletableFuture +import java.util.concurrent.atomic.AtomicBoolean + +internal const val X509_API_BASE_URL = "https://mtls.api.openai.com/v1" + +internal class X509ClientConfiguration +private constructor( + private val identity: X509WorkloadIdentity, + private val bindTransport: (Timeout) -> BoundX509Transport, +) { + companion object { + @JvmSynthetic + internal fun create( + identity: X509WorkloadIdentity, + bindTransport: (Timeout) -> BoundX509Transport, + ) = X509ClientConfiguration(identity, bindTransport) + } + + @JvmSynthetic + fun reserve(clientOptions: ClientOptions.Builder) { + clientOptions.fixedBearerAuthentication(X509_API_BASE_URL) + } + + @JvmSynthetic + fun buildClientOptions(clientOptions: ClientOptions.Builder): ClientOptions { + val transport = bindTransport(clientOptions.timeout()) + return try { + clientOptions + .fixedBearerTransport( + transport.apiClient, + X509RequestAuthenticator(identity, transport.exchangeClient), + ) + .build() + } catch (error: Throwable) { + try { + transport.close() + } catch (closeError: Throwable) { + if (closeError !== error) { + error.addSuppressed(closeError) + } + } + throw error + } + } +} + +/** Owns the exchange client and adds a freshly exchanged bearer to API requests. */ +private class X509RequestAuthenticator( + identity: X509WorkloadIdentity, + private val exchangeClient: OkHttpClient, +) : HttpRequestAuthenticator { + private val tokenExchange = X509TokenExchange(identity, exchangeClient) + private val closed = AtomicBoolean() + + override fun authenticate(request: HttpRequest): HttpRequest { + validateRequest(request) + return authenticated(request, tokenExchange.execute()) + } + + override fun authenticateAsync(request: HttpRequest): CompletableFuture { + try { + validateRequest(request) + } catch (error: Throwable) { + return CompletableFuture().also { it.completeExceptionally(error) } + } + + val exchangeFuture = tokenExchange.executeAsync() + val result = CompletableFuture() + exchangeFuture.whenComplete { token, error -> + if (error != null) { + result.completeExceptionally(error) + } else if (token == null) { + result.completeExceptionally( + IllegalStateException("X.509 token exchange completed without a token") + ) + } else { + try { + result.complete(authenticated(request, token)) + } catch (authenticationError: Throwable) { + result.completeExceptionally(authenticationError) + } + } + } + result.whenComplete { _, error -> + if (error is CancellationException) { + exchangeFuture.cancel(true) + } + } + return result + } + + override fun close() { + if (closed.compareAndSet(false, true)) { + exchangeClient.close() + } + } + + private fun authenticated(request: HttpRequest, token: X509AccessToken): HttpRequest = + request.toBuilder().replaceHeaders("Authorization", "Bearer ${token.value}").build() + + private fun validateRequest(request: HttpRequest) { + require(request.baseUrl == X509_API_BASE_URL) { + "X.509 workload identity is restricted to $X509_API_BASE_URL" + } + request.headers.names().forEach { name -> + val normalized = name.trim().lowercase(Locale.ROOT).replace('_', '-') + require(normalized !in FORBIDDEN_HEADERS) { + "Header $name cannot be configured with X.509 workload identity" + } + } + } + + private companion object { + val FORBIDDEN_HEADERS = + setOf( + "authorization", + "api-key", + "x-api-key", + "proxy-authorization", + "cookie", + "host", + ":authority", + "openai-organization", + "openai-project", + ) + } +} diff --git a/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/X509Transport.kt b/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/X509Transport.kt index b0d0cc247..06835d008 100644 --- a/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/X509Transport.kt +++ b/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/X509Transport.kt @@ -114,7 +114,7 @@ private constructor( val exchangeClient = client(exchangeProxy) return try { - BoundX509Transport(exchangeClient, client(apiProxy)) + BoundX509Transport.create(exchangeClient, client(apiProxy)) } catch (error: Throwable) { try { exchangeClient.close() @@ -128,8 +128,14 @@ private constructor( } } -internal class BoundX509Transport(val exchangeClient: OkHttpClient, val apiClient: OkHttpClient) : - AutoCloseable { +internal class BoundX509Transport +private constructor(val exchangeClient: OkHttpClient, val apiClient: OkHttpClient) : AutoCloseable { + + companion object { + @JvmSynthetic + internal fun create(exchangeClient: OkHttpClient, apiClient: OkHttpClient) = + BoundX509Transport(exchangeClient, apiClient) + } override fun close() { apiClient.use { exchangeClient.close() } diff --git a/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/OkHttpClientTest.kt b/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/OkHttpClientTest.kt index ed1c83c02..f22297b9e 100644 --- a/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/OkHttpClientTest.kt +++ b/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/OkHttpClientTest.kt @@ -9,6 +9,13 @@ 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 +import okhttp3.mockwebserver.MockResponse +import okhttp3.mockwebserver.MockWebServer +import okhttp3.mockwebserver.SocketPolicy import org.assertj.core.api.Assertions.assertThat import org.junit.jupiter.api.BeforeEach import org.junit.jupiter.api.Test @@ -68,6 +75,54 @@ internal class OkHttpClientTest { assertThat(future.get()).isSameAs(response) assertThat(response.closed).isFalse() } + + @Test + fun close_cancelsSyncAndAsyncResponseBodyReadsAfterHeaders() { + listOf(false, true).forEach { async -> + val server = MockWebServer() + val client = OkHttpClient.builder().build() + val executor = Executors.newCachedThreadPool() + try { + server.start() + server.enqueue( + MockResponse() + .setBody("partial") + .setHeader("Content-Length", 100) + .setSocketPolicy(SocketPolicy.KEEP_OPEN) + ) + val request = + HttpRequest.builder() + .method(HttpMethod.GET) + .baseUrl(server.url("/").toString()) + .build() + val response = + if (async) client.executeAsync(request).get(5, TimeUnit.SECONDS) + else + executor + .submit { client.execute(request) } + .get(5, TimeUnit.SECONDS) + val readStarted = CountDownLatch(1) + val readFuture = + executor.submit { + readStarted.countDown() + response.body().readBytes().size + } + assertThat(readStarted.await(5, TimeUnit.SECONDS)).isTrue() + + client.close() + + try { + readFuture.get(5, TimeUnit.SECONDS) + } catch (_: ExecutionException) {} + assertThat(readFuture.isDone).isTrue() + response.close() + } finally { + client.close() + executor.shutdownNow() + server.close() + } + } + } } private class TrackingHttpResponse : HttpResponse { diff --git a/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/OpenAIOkHttpClientX509Test.kt b/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/OpenAIOkHttpClientX509Test.kt new file mode 100644 index 000000000..7ee728dd6 --- /dev/null +++ b/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/OpenAIOkHttpClientX509Test.kt @@ -0,0 +1,415 @@ +package com.openai.client.okhttp + +import com.fasterxml.jackson.databind.ObjectMapper +import com.openai.auth.X509WorkloadIdentity +import com.openai.credential.BearerTokenCredential +import com.openai.models.files.FileListParams +import java.net.Proxy +import java.security.cert.X509Certificate +import java.time.Duration +import java.util.concurrent.ExecutionException +import java.util.concurrent.Executors +import java.util.concurrent.TimeUnit +import okhttp3.mockwebserver.MockResponse +import okhttp3.mockwebserver.RecordedRequest +import okhttp3.mockwebserver.SocketPolicy +import org.assertj.core.api.Assertions.assertThat +import org.assertj.core.api.Assertions.assertThatThrownBy +import org.junit.jupiter.api.Test + +internal class OpenAIOkHttpClientX509Test { + private val jsonMapper = ObjectMapper() + + @Test + fun syncClientUsesExactExchangeAndApiWireContract() { + Fixture().use { fixture -> + fixture.enqueueSuccess() + val client = + fixture + .syncBuilder() + .putHeader("X-Application-Test", "api-only") + .build() + .withOptions { options -> options.timeout(Duration.ofSeconds(5)) } + + try { + assertThat(client.files().list().data()).isEmpty() + } finally { + client.close() + } + + assertSuccessfulWireContract( + fixture, + expectedUserAgent = "OpenAIClientImpl/Java", + expectedApplicationHeader = "api-only", + ) + } + } + + @Test + fun asyncClientUsesExactExchangeAndApiWireContract() { + Fixture().use { fixture -> + fixture.enqueueSuccess() + val client = fixture.asyncBuilder().build() + + try { + assertThat(client.files().list().get(10, TimeUnit.SECONDS).data()).isEmpty() + } finally { + client.close() + } + + assertSuccessfulWireContract(fixture, expectedUserAgent = "OpenAIClientAsyncImpl/Java") + } + } + + @Test + fun onlyPublicJavaConstructionPathIsX509BuilderFactory() { + listOf(OpenAIOkHttpClient.Builder::class.java, OpenAIOkHttpClientAsync.Builder::class.java) + .forEach { builderClass -> + assertThat(builderClass.constructors).noneMatch { constructor -> + !constructor.isSynthetic && + constructor.parameterTypes.contains(X509ClientConfiguration::class.java) + } + } + assertThat( + X509ClientConfiguration::class.java.constructors.filterNot { constructor -> + constructor.isSynthetic + } + ) + .isEmpty() + assertThat( + BoundX509Transport::class.java.constructors.filterNot { constructor -> + constructor.isSynthetic + } + ) + .isEmpty() + assertThat( + X509ClientConfiguration::class.java.declaredMethods.filterNot { method -> + method.isSynthetic + } + ) + .isEmpty() + assertThat( + OpenAIOkHttpClient::class + .java + .getMethod( + "x509Builder", + X509WorkloadIdentity::class.java, + X509Transport::class.java, + ) + ) + .isNotNull() + assertThat( + OpenAIOkHttpClientAsync::class + .java + .getMethod( + "x509Builder", + X509WorkloadIdentity::class.java, + X509Transport::class.java, + ) + ) + .isNotNull() + } + + @Test + fun fixedModeRejectsCompetingConfigurationBeforeNetworkUse() { + Fixture().use { fixture -> + val syncMutations = + listOf Unit>>( + "baseUrl" to { it.baseUrl("https://example.test/v1") }, + "apiKey" to { it.apiKey("test-api-key") }, + "adminApiKey" to { it.adminApiKey("test-admin-key") }, + "credential" to { it.credential(BearerTokenCredential.create("test-token")) }, + "organization" to { it.organization("org_test") }, + "project" to { it.project("proj_test") }, + "fromEnv" to { it.fromEnv() }, + "proxy" to { it.proxy(Proxy.NO_PROXY) }, + "followRedirects" to { it.followRedirects(false) }, + ) + syncMutations.forEach { (name, mutate) -> + assertThatThrownBy { mutate(fixture.syncBuilder()) } + .isInstanceOf(IllegalArgumentException::class.java) + .hasMessageContaining(name) + } + + assertThatThrownBy { fixture.asyncBuilder().baseUrl("https://example.test/v1") } + .isInstanceOf(IllegalArgumentException::class.java) + .hasMessageContaining("baseUrl") + assertThatThrownBy { fixture.asyncBuilder().proxy(Proxy.NO_PROXY) } + .isInstanceOf(IllegalArgumentException::class.java) + .hasMessageContaining("proxy") + assertThat(fixture.authPeer.server.requestCount).isZero() + assertThat(fixture.apiPeer.server.requestCount).isZero() + } + } + + @Test + fun withOptionsCannotMoveBearerToAnotherEndpoint() { + Fixture().use { fixture -> + val client = fixture.syncBuilder().build() + try { + assertThatThrownBy { + client.withOptions { options -> + options.baseUrl("https://attacker.example/v1") + } + } + .isInstanceOf(IllegalArgumentException::class.java) + .hasMessageContaining("baseUrl") + assertThatThrownBy { + client.files().withOptions { options -> + options.apiKey("replacement-secret") + } + } + .isInstanceOf(IllegalArgumentException::class.java) + .hasMessageContaining("apiKey") + } finally { + client.close() + } + assertThat(fixture.authPeer.server.requestCount).isZero() + assertThat(fixture.apiPeer.server.requestCount).isZero() + } + } + + @Test + fun callerSuppliedCredentialHeaderIsRejectedBeforeTokenExchange() { + Fixture().use { fixture -> + val client = fixture.syncBuilder().putHeader("Authorization", "Bearer secret").build() + try { + assertThatThrownBy { client.files().list() } + .isInstanceOf(IllegalArgumentException::class.java) + .hasMessageContaining("Authorization") + .hasMessageNotContaining("secret") + } finally { + client.close() + } + assertThat(fixture.authPeer.server.requestCount).isZero() + assertThat(fixture.apiPeer.server.requestCount).isZero() + } + } + + @Test + fun requestLevelCredentialHeaderIsRejectedBeforeTokenExchange() { + Fixture().use { fixture -> + val client = fixture.syncBuilder().build() + val params = + FileListParams.builder() + .putAdditionalHeader("api_key", "request-level-secret") + .build() + try { + assertThatThrownBy { client.files().list(params) } + .isInstanceOf(IllegalArgumentException::class.java) + .hasMessageContaining("api_key") + .hasMessageNotContaining("request-level-secret") + } finally { + client.close() + } + assertThat(fixture.authPeer.server.requestCount).isZero() + assertThat(fixture.apiPeer.server.requestCount).isZero() + } + } + + @Test + fun adminOnlyRouteFailsLocallyWithoutExchangingToken() { + Fixture().use { fixture -> + val client = fixture.syncBuilder().build() + try { + assertThatThrownBy { client.admin().organization().auditLogs().list() } + .isInstanceOf(IllegalStateException::class.java) + .hasMessageContaining("requires adminApiKey") + } finally { + client.close() + } + assertThat(fixture.authPeer.server.requestCount).isZero() + assertThat(fixture.apiPeer.server.requestCount).isZero() + } + } + + @Test + fun closingAsyncClientCancelsBlockedExchangeBeforeApiDispatch() { + Fixture().use { fixture -> + fixture.authPeer.enqueue(MockResponse().setSocketPolicy(SocketPolicy.NO_RESPONSE)) + val client = fixture.asyncBuilder().timeout(Duration.ofSeconds(30)).build() + val future = client.files().list() + + assertConnectAuthority(fixture.authPeer.takeRequest(), AUTH_AUTHORITY) + assertThat(fixture.authPeer.takeRequest().path).isEqualTo("/oauth/token") + client.close() + + assertThatThrownBy { future.get(5, TimeUnit.SECONDS) } + .isInstanceOf(ExecutionException::class.java) + assertThat(fixture.apiPeer.server.requestCount).isZero() + } + } + + @Test + fun closingSyncClientCancelsBlockedExchangeBeforeApiDispatch() { + Fixture().use { fixture -> + fixture.authPeer.enqueue(MockResponse().setSocketPolicy(SocketPolicy.NO_RESPONSE)) + val client = fixture.syncBuilder().timeout(Duration.ofSeconds(30)).build() + val executor = Executors.newSingleThreadExecutor() + val future = executor.submit { client.files().list() } + + try { + assertConnectAuthority(fixture.authPeer.takeRequest(), AUTH_AUTHORITY) + assertThat(fixture.authPeer.takeRequest().path).isEqualTo("/oauth/token") + client.close() + + assertThatThrownBy { future.get(5, TimeUnit.SECONDS) } + .isInstanceOf(ExecutionException::class.java) + assertThat(fixture.apiPeer.server.requestCount).isZero() + } finally { + client.close() + executor.shutdownNow() + } + } + } + + @Test + fun closingAsyncClientCancelsBlockedApiDispatch() { + Fixture().use { fixture -> + fixture.authPeer.enqueue( + MockResponse().setHeader("Content-Type", "application/json").setBody(TOKEN_RESPONSE) + ) + fixture.apiPeer.enqueue(MockResponse().setSocketPolicy(SocketPolicy.NO_RESPONSE)) + val client = fixture.asyncBuilder().timeout(Duration.ofSeconds(30)).build() + val future = client.files().list() + + assertConnectAuthority(fixture.authPeer.takeRequest(), AUTH_AUTHORITY) + assertThat(fixture.authPeer.takeRequest().path).isEqualTo("/oauth/token") + assertConnectAuthority(fixture.apiPeer.takeRequest(), API_AUTHORITY) + assertThat(fixture.apiPeer.takeRequest().path).isEqualTo("/v1/files") + client.close() + + assertThatThrownBy { future.get(5, TimeUnit.SECONDS) } + .isInstanceOf(ExecutionException::class.java) + assertThat(fixture.authPeer.server.requestCount).isEqualTo(2) + assertThat(fixture.apiPeer.server.requestCount).isEqualTo(2) + } + } + + private fun assertSuccessfulWireContract( + fixture: Fixture, + expectedUserAgent: String, + expectedApplicationHeader: String? = null, + ) { + val authConnect = fixture.authPeer.takeRequest() + val exchangeRequest = fixture.authPeer.takeRequest() + val apiConnect = fixture.apiPeer.takeRequest() + val apiRequest = fixture.apiPeer.takeRequest() + + assertConnectAuthority(authConnect, AUTH_AUTHORITY) + assertConnectAuthority(apiConnect, API_AUTHORITY) + assertThat(exchangeRequest.method).isEqualTo("POST") + assertThat(exchangeRequest.path).isEqualTo("/oauth/token") + assertThat(exchangeRequest.getHeader("Authorization")).isNull() + assertThat(exchangeRequest.getHeader("X-Application-Test")).isNull() + assertThat(jsonMapper.readTree(exchangeRequest.body.readUtf8())) + .isEqualTo(jsonMapper.readTree(TOKEN_REQUEST)) + assertThat(apiRequest.method).isEqualTo("GET") + assertThat(apiRequest.path).isEqualTo("/v1/files") + assertThat(apiRequest.headers.values("Authorization")) + .containsExactly("Bearer $ACCESS_TOKEN") + assertThat(apiRequest.getHeader("api-key")).isNull() + assertThat(apiRequest.getHeader("Cookie")).isNull() + assertThat(apiRequest.getHeader("OpenAI-Organization")).isNull() + assertThat(apiRequest.getHeader("OpenAI-Project")).isNull() + assertThat(apiRequest.getHeader("X-Application-Test")).isEqualTo(expectedApplicationHeader) + assertThat(apiRequest.getHeader("User-Agent")).startsWith(expectedUserAgent) + assertPresentedIdentity(exchangeRequest, fixture.identity.leaf.certificate) + assertPresentedIdentity(apiRequest, fixture.identity.leaf.certificate) + assertThat(exchangeRequest.handshake!!.peerCertificates) + .isEqualTo(apiRequest.handshake!!.peerCertificates) + assertThat(fixture.authPeer.requestedServerNames).containsExactly(AUTH_HOST) + assertThat(fixture.apiPeer.requestedServerNames).containsExactly(API_HOST) + } + + private fun assertConnectAuthority(request: RecordedRequest, authority: String) { + assertThat(request.requestLine).isEqualTo("CONNECT $authority HTTP/1.1") + } + + private fun assertPresentedIdentity(request: RecordedRequest, expected: X509Certificate) { + assertThat(requireNotNull(request.handshake).peerCertificates.first()).isEqualTo(expected) + } + + private class Fixture : AutoCloseable { + val identity = X509TestIdentity.create("SDK X.509 identity") + val authPeer = X509TestPeer(AUTH_HOST, identity.root.certificate) + val apiPeer = X509TestPeer(API_HOST, identity.root.certificate) + private val workloadIdentity = + X509WorkloadIdentity.builder() + .identityProviderId("idp_test") + .serviceAccountId("svc_acct_test") + .build() + private val transport = + X509Transport.builder() + .keyManager(x509TestKeyManager(mapOf(CERTIFICATE_ALIAS to identity))) + .certificateAlias(CERTIFICATE_ALIAS) + .trustManager( + identity + .clientHandshakeCertificates( + listOf(authPeer.serverRootCertificate, apiPeer.serverRootCertificate) + ) + .trustManager + ) + .build() + + fun syncBuilder(): OpenAIOkHttpClient.Builder = + OpenAIOkHttpClient.x509BuilderForTest( + workloadIdentity, + transport, + authPeer.proxy, + apiPeer.proxy, + ) + + fun asyncBuilder(): OpenAIOkHttpClientAsync.Builder = + OpenAIOkHttpClientAsync.x509BuilderForTest( + workloadIdentity, + transport, + authPeer.proxy, + apiPeer.proxy, + ) + + fun enqueueSuccess() { + authPeer.enqueue( + MockResponse().setHeader("Content-Type", "application/json").setBody(TOKEN_RESPONSE) + ) + apiPeer.enqueue( + MockResponse().setHeader("Content-Type", "application/json").setBody(FILES_RESPONSE) + ) + } + + override fun close() { + apiPeer.use { authPeer.close() } + } + } + + private companion object { + const val AUTH_HOST = "mtls.auth.openai.com" + const val API_HOST = "mtls.api.openai.com" + const val AUTH_AUTHORITY = "$AUTH_HOST:443" + const val API_AUTHORITY = "$API_HOST:443" + const val ACCESS_TOKEN = "test-x509-access-token" + const val CERTIFICATE_ALIAS = "selected" + const val ACCESS_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:access_token" + const val FILES_RESPONSE = """{"object":"list","data":[]}""" + val TOKEN_REQUEST = + """ + { + "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange", + "subject_token_type": "urn:openai:params:oauth:token-type:x509", + "identity_provider_id": "idp_test", + "service_account_id": "svc_acct_test" + } + """ + .trimIndent() + val TOKEN_RESPONSE = + """ + { + "access_token": "$ACCESS_TOKEN", + "issued_token_type": "$ACCESS_TOKEN_TYPE", + "token_type": "Bearer", + "expires_in": 3600 + } + """ + .trimIndent() + } +} diff --git a/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/X509WireTestInfrastructure.kt b/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/X509WireTestInfrastructure.kt index fa6932d00..57ee080d8 100644 --- a/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/X509WireTestInfrastructure.kt +++ b/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/X509WireTestInfrastructure.kt @@ -5,7 +5,9 @@ import java.net.Socket import java.security.KeyStore import java.security.SecureRandom import java.security.cert.X509Certificate +import java.time.Duration import java.util.concurrent.CopyOnWriteArrayList +import java.util.concurrent.TimeUnit import javax.net.ssl.ExtendedSSLSession import javax.net.ssl.KeyManagerFactory import javax.net.ssl.SNIHostName @@ -97,7 +99,10 @@ internal class X509TestPeer(val authority: String, trustedClientRoot: X509Certif server.enqueue(response) } - fun takeRequest(): RecordedRequest = server.takeRequest() + fun takeRequest(timeout: Duration = Duration.ofSeconds(5)): RecordedRequest = + requireNotNull(server.takeRequest(timeout.toMillis(), TimeUnit.MILLISECONDS)) { + "No request received by $authority within $timeout" + } override fun close() { server.close() diff --git a/openai-java-core/src/main/kotlin/com/openai/core/ClientOptions.kt b/openai-java-core/src/main/kotlin/com/openai/core/ClientOptions.kt index 277954de3..85c70d186 100644 --- a/openai-java-core/src/main/kotlin/com/openai/core/ClientOptions.kt +++ b/openai-java-core/src/main/kotlin/com/openai/core/ClientOptions.kt @@ -14,7 +14,6 @@ import com.openai.core.http.HttpClient import com.openai.core.http.HttpRequestAuthenticator import com.openai.core.http.LoggingHttpClient import com.openai.core.http.PhantomReachableClosingHttpClient -import com.openai.core.http.PhantomReachableClosingHttpRequestAuthenticator import com.openai.core.http.QueryParams import com.openai.core.http.RetryingHttpClient import com.openai.core.http.WorkloadIdentityHttpClient @@ -43,7 +42,7 @@ private constructor( * This class takes ownership of the client and closes it when closed. */ @get:JvmName("httpClient") val httpClient: HttpClient, - private val httpRequestAuthenticator: HttpRequestAuthenticator?, + private val requestAuthentication: RequestAuthentication, /** * Whether to throw an exception if any of the Jackson versions detected at runtime are * incompatible with the SDK's minimum supported Jackson version (2.13.4). @@ -194,7 +193,7 @@ private constructor( class Builder internal constructor() { private var httpClient: HttpClient? = null - private var httpRequestAuthenticator: HttpRequestAuthenticator? = null + private var requestAuthentication: RequestAuthentication = RequestAuthentication.None private var checkJacksonVersionCompatibility: Boolean = true private var jsonMapper: JsonMapper = jsonMapper() private var streamHandlerExecutor: Executor? = null @@ -225,7 +224,7 @@ private constructor( @JvmSynthetic internal fun from(clientOptions: ClientOptions) = apply { httpClient = clientOptions.originalHttpClient - httpRequestAuthenticator = clientOptions.httpRequestAuthenticator + requestAuthentication = clientOptions.requestAuthentication checkJacksonVersionCompatibility = clientOptions.checkJacksonVersionCompatibility jsonMapper = clientOptions.jsonMapper streamHandlerExecutor = clientOptions.streamHandlerExecutor @@ -264,6 +263,7 @@ private constructor( * This class takes ownership of the client and closes it when closed. */ fun httpClient(httpClient: HttpClient) = apply { + requireNoFixedBearerAuthentication("httpClient") this.httpClient = PhantomReachableClosingHttpClient(httpClient) } @@ -275,9 +275,52 @@ private constructor( */ @JvmSynthetic fun httpRequestAuthenticator(httpRequestAuthenticator: HttpRequestAuthenticator?) = apply { - this.httpRequestAuthenticator = - if (httpRequestAuthenticator == null) null - else PhantomReachableClosingHttpRequestAuthenticator(httpRequestAuthenticator) + requireNoFixedBearerAuthentication("httpRequestAuthenticator") + requestAuthentication = + httpRequestAuthenticator?.let(RequestAuthentication.Provider::create) + ?: RequestAuthentication.None + } + + /** Reserves a fixed-origin, bearer-only authentication mode for an SDK integration. */ + @JvmSynthetic + fun fixedBearerAuthentication(fixedBaseUrl: String) = apply { + require(requestAuthentication === RequestAuthentication.None) { + "Fixed bearer authentication is already set" + } + require( + httpClient == null && + baseUrl == null && + !dataResidencySelected && + credential == null && + adminApiKey.isNullOrEmpty() && + workloadIdentity == null && + azureServiceVersion == null && + azureUrlPathMode == AzureUrlPathMode.AUTO && + organization == null && + project == null + ) { + "Fixed bearer authentication cannot be combined with existing client, endpoint, credential, provider, organization, or project configuration" + } + require(fixedBaseUrl.isNotBlank()) { "fixedBaseUrl must not be blank" } + requestAuthentication = RequestAuthentication.FixedBearerReserved(fixedBaseUrl) + baseUrl = fixedBaseUrl + explicitBaseUrl = true + } + + /** Installs the transport owned by a previously reserved fixed-bearer integration. */ + @JvmSynthetic + fun fixedBearerTransport( + httpClient: HttpClient, + httpRequestAuthenticator: HttpRequestAuthenticator, + ) = apply { + val reserved = requestAuthentication as? RequestAuthentication.FixedBearerReserved + checkNotNull(reserved) { "Fixed bearer authentication must be set first" } + this.httpClient = PhantomReachableClosingHttpClient(httpClient) + requestAuthentication = + RequestAuthentication.FixedBearerInstalled.create( + reserved.fixedBearerBaseUrl, + httpRequestAuthenticator, + ) } /** @@ -339,6 +382,7 @@ private constructor( * Defaults to the production environment: `https://api.openai.com/v1`. */ fun baseUrl(baseUrl: String?) = apply { + requireNoFixedBearerAuthentication("baseUrl") require(!explicitDataResidency) { "baseUrl and dataResidency are mutually exclusive" } this.baseUrl = baseUrl dataResidencySelected = false @@ -356,6 +400,7 @@ private constructor( * leaves the endpoint unchanged. Availability is determined by the API. */ fun dataResidency(dataResidency: DataResidency?) = apply { + requireNoFixedBearerAuthentication("dataResidency") if (dataResidency != null) { require(!explicitBaseUrl) { "baseUrl and dataResidency are mutually exclusive" } baseUrl = dataResidency.baseUrl @@ -425,6 +470,7 @@ private constructor( fun logLevel(logLevel: LogLevel) = apply { this.logLevel = logLevel } fun apiKey(apiKey: String?) = apply { + requireNoFixedBearerAuthentication("apiKey") this.apiKey = apiKey this.credential = apiKey?.let { BearerTokenCredential.create(it) } } @@ -432,30 +478,42 @@ private constructor( /** Alias for calling [Builder.apiKey] with `apiKey.orElse(null)`. */ fun apiKey(apiKey: Optional) = apiKey(apiKey.getOrNull()) - fun adminApiKey(adminApiKey: String?) = apply { this.adminApiKey = adminApiKey } + fun adminApiKey(adminApiKey: String?) = apply { + requireNoFixedBearerAuthentication("adminApiKey") + this.adminApiKey = adminApiKey + } /** Alias for calling [Builder.adminApiKey] with `adminApiKey.orElse(null)`. */ fun adminApiKey(adminApiKey: Optional) = adminApiKey(adminApiKey.getOrNull()) fun credential(credential: Credential) = apply { + requireNoFixedBearerAuthentication("credential") this.apiKey = null this.credential = credential } fun azureServiceVersion(azureServiceVersion: AzureOpenAIServiceVersion) = apply { + requireNoFixedBearerAuthentication("azureServiceVersion") this.azureServiceVersion = azureServiceVersion } fun azureUrlPathMode(azureUrlPathMode: AzureUrlPathMode) = apply { + requireNoFixedBearerAuthentication("azureUrlPathMode") this.azureUrlPathMode = azureUrlPathMode } - fun organization(organization: String?) = apply { this.organization = organization } + fun organization(organization: String?) = apply { + requireNoFixedBearerAuthentication("organization") + this.organization = organization + } /** Alias for calling [Builder.organization] with `organization.orElse(null)`. */ fun organization(organization: Optional) = organization(organization.getOrNull()) - fun project(project: String?) = apply { this.project = project } + fun project(project: String?) = apply { + requireNoFixedBearerAuthentication("project") + this.project = project + } /** Alias for calling [Builder.project] with `project.orElse(null)`. */ fun project(project: Optional) = project(project.getOrNull()) @@ -467,6 +525,7 @@ private constructor( webhookSecret(webhookSecret.getOrNull()) fun workloadIdentity(workloadIdentity: WorkloadIdentity?) = apply { + requireNoFixedBearerAuthentication("workloadIdentity") this.workloadIdentity = workloadIdentity } @@ -556,11 +615,30 @@ private constructor( fun timeout(): Timeout = timeout + private fun requireNoFixedBearerAuthentication(option: String) { + require(requestAuthentication.fixedBearerBaseUrl == null) { + "$option cannot be configured with fixed bearer authentication" + } + } + private fun effectiveCredential( httpClient: HttpClient, jsonMapper: JsonMapper, ): Credential { - if (httpRequestAuthenticator != null) { + if (requestAuthentication is RequestAuthentication.FixedBearerReserved) { + throw IllegalStateException( + "Fixed bearer authentication transport is not configured" + ) + } + if (requestAuthentication is RequestAuthentication.FixedBearerInstalled) { + check( + credential == null && workloadIdentity == null && adminApiKey.isNullOrEmpty() + ) { + "Fixed bearer authentication cannot be combined with other credentials" + } + return HttpRequestAuthenticatorCredential + } + if (requestAuthentication is RequestAuthentication.Provider) { if ( credential != null || workloadIdentity != null || !adminApiKey.isNullOrEmpty() ) { @@ -615,6 +693,7 @@ private constructor( * System properties take precedence over environment variables. */ fun fromEnv() = apply { + requireNoFixedBearerAuthentication("fromEnv") logLevel(LogLevel.fromEnv()) (System.getProperty("openai.baseUrl") ?: System.getenv("OPENAI_BASE_URL"))?.let { if (!dataResidencySelected) { @@ -678,10 +757,16 @@ private constructor( * @throws IllegalStateException if any required field is unset. */ fun build(): ClientOptions { + require( + requestAuthentication.fixedBearerBaseUrl == null || + baseUrl == requestAuthentication.fixedBearerBaseUrl + ) { + "Fixed bearer authentication base URL cannot be changed" + } require( !dataResidencySelected || (!inheritedAzureEndpoint && - httpRequestAuthenticator == null && + requestAuthentication === RequestAuthentication.None && credential !is AzureApiKeyCredential && azureServiceVersion == null && azureUrlPathMode == AzureUrlPathMode.AUTO) @@ -750,7 +835,7 @@ private constructor( (credential as? WorkloadIdentityCredential)?.getAuth() val loggingDelegate = - if (httpRequestAuthenticator != null) httpClient + if (requestAuthentication.authenticator != null) httpClient else WorkloadIdentityHttpClient( delegate = httpClient, @@ -765,7 +850,7 @@ private constructor( .build() val perAttemptHttpClient = - httpRequestAuthenticator?.let { authenticator -> + requestAuthentication.authenticator?.let { authenticator -> AuthenticatingHttpClient( delegate = loggingHttpClient, authenticator = authenticator, @@ -783,7 +868,7 @@ private constructor( return ClientOptions( httpClient, wrappedHttpClient, - httpRequestAuthenticator, + requestAuthentication, checkJacksonVersionCompatibility, jsonMapper, streamHandlerExecutor, @@ -828,9 +913,7 @@ private constructor( @JvmSynthetic internal fun securityHeaders(security: SecurityOptions): Headers { val headers = Headers.builder() - var isSatisfied = - credential === HttpRequestAuthenticatorCredential && - (security.bearerAuth || security.adminApiKeyAuth) + var isSatisfied = requestAuthentication.satisfies(security) if (security.bearerAuth) { when { @@ -848,9 +931,6 @@ private constructor( credential is WorkloadIdentityCredential -> { isSatisfied = true } - credential === HttpRequestAuthenticatorCredential -> { - isSatisfied = true - } } } if (security.adminApiKeyAuth) { diff --git a/openai-java-core/src/main/kotlin/com/openai/core/RequestAuthentication.kt b/openai-java-core/src/main/kotlin/com/openai/core/RequestAuthentication.kt new file mode 100644 index 000000000..f376b33b2 --- /dev/null +++ b/openai-java-core/src/main/kotlin/com/openai/core/RequestAuthentication.kt @@ -0,0 +1,59 @@ +package com.openai.core + +import com.openai.core.http.HttpRequestAuthenticator +import com.openai.core.http.PhantomReachableClosingHttpRequestAuthenticator + +/** Authentication that is applied to the final HTTP request rather than represented by a key. */ +internal sealed interface RequestAuthentication { + val authenticator: HttpRequestAuthenticator? + + val fixedBearerBaseUrl: String? + get() = null + + fun satisfies(security: SecurityOptions): Boolean + + object None : RequestAuthentication { + override val authenticator: HttpRequestAuthenticator? = null + + override fun satisfies(security: SecurityOptions): Boolean = false + } + + class Provider private constructor(override val authenticator: HttpRequestAuthenticator) : + RequestAuthentication { + + override fun satisfies(security: SecurityOptions): Boolean = + security.bearerAuth || security.adminApiKeyAuth + + companion object { + fun create(authenticator: HttpRequestAuthenticator): Provider = + Provider(PhantomReachableClosingHttpRequestAuthenticator(authenticator)) + } + } + + class FixedBearerReserved(override val fixedBearerBaseUrl: String) : RequestAuthentication { + override val authenticator: HttpRequestAuthenticator? = null + + override fun satisfies(security: SecurityOptions): Boolean = false + } + + class FixedBearerInstalled + private constructor( + override val fixedBearerBaseUrl: String, + override val authenticator: HttpRequestAuthenticator, + ) : RequestAuthentication { + + override fun satisfies(security: SecurityOptions): Boolean = + security.bearerAuth && !security.adminApiKeyAuth + + companion object { + fun create( + fixedBearerBaseUrl: String, + authenticator: HttpRequestAuthenticator, + ): FixedBearerInstalled = + FixedBearerInstalled( + fixedBearerBaseUrl, + PhantomReachableClosingHttpRequestAuthenticator(authenticator), + ) + } + } +} diff --git a/openai-java-core/src/test/kotlin/com/openai/core/ClientOptionsTest.kt b/openai-java-core/src/test/kotlin/com/openai/core/ClientOptionsTest.kt index 2817c98ef..81c2416b8 100644 --- a/openai-java-core/src/test/kotlin/com/openai/core/ClientOptionsTest.kt +++ b/openai-java-core/src/test/kotlin/com/openai/core/ClientOptionsTest.kt @@ -179,6 +179,73 @@ internal class ClientOptionsTest { assertThat(thrown.message).contains("Provider authentication cannot be combined") } + @Test + fun build_withFixedBearerAuthentication_satisfiesOnlyBearerAndSurvivesCloning() { + val authenticator = + object : HttpRequestAuthenticator { + override fun authenticate(request: HttpRequest): HttpRequest = request + } + val clientOptions = + ClientOptions.builder() + .fixedBearerAuthentication("https://mtls.example.test/v1") + .fixedBearerTransport(httpClient, authenticator) + .build() + .toBuilder() + .build() + + assertThat(clientOptions.baseUrl()).isEqualTo("https://mtls.example.test/v1") + assertThat( + clientOptions.securityHeaders(SecurityOptions.builder().bearerAuth(true).build()) + ) + .isEqualTo(com.openai.core.http.Headers.builder().build()) + val thrown = + assertThrows { + clientOptions.securityHeaders( + SecurityOptions.builder().adminApiKeyAuth(true).build() + ) + } + assertThat(thrown.message).contains("requires adminApiKey") + } + + @Test + fun fixedBearerAuthentication_rejectsCompetingConfigurationInEitherOrder() { + val mutations = + listOf Unit>>( + "httpClient" to { it.httpClient(httpClient) }, + "httpRequestAuthenticator" to + { + it.httpRequestAuthenticator( + object : HttpRequestAuthenticator { + override fun authenticate(request: HttpRequest): HttpRequest = + request + } + ) + }, + "baseUrl" to { it.baseUrl("https://example.test/v1") }, + "apiKey" to { it.apiKey("test-api-key") }, + "adminApiKey" to { it.adminApiKey("test-admin-key") }, + "credential" to { it.credential(BearerTokenCredential.create("test-token")) }, + "organization" to { it.organization("org_test") }, + "project" to { it.project("proj_test") }, + "fromEnv" to { it.fromEnv() }, + ) + + mutations.forEach { (name, mutate) -> + val builder = + ClientOptions.builder().fixedBearerAuthentication("https://mtls.example.test/v1") + val thrown = assertThrows { mutate(builder) } + assertThat(thrown.message).contains(name) + } + + val thrown = + assertThrows { + ClientOptions.builder() + .apiKey("test-api-key") + .fixedBearerAuthentication("https://mtls.example.test/v1") + } + assertThat(thrown.message).contains("cannot be combined") + } + @Test fun putHeader_canOverwriteDefaultHeader() { val clientOptions = From a21a0b2c1c83b0fab582fadc857657b849d064b1 Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Tue, 25 Aug 2026 20:07:26 +0000 Subject: [PATCH 2/5] feat(auth): harden X.509 request lifecycle --- .github/workflows/x509-live-smoke.yml | 69 ++ README.md | 39 + .../openai/gradle/CoreCompilationShards.kt | 4 + .../client/okhttp/X509ClientIntegration.kt | 332 +++++- .../openai/client/okhttp/X509TokenExchange.kt | 11 +- .../okhttp/OpenAIOkHttpClientX509Test.kt | 256 ++++- .../okhttp/X509AttemptAuthenticatorTest.kt | 411 ++++++++ .../openai/client/OpenAIClientAsyncImpl.kt | 8 +- .../com/openai/client/OpenAIClientImpl.kt | 8 +- .../core/CancellationPropagatingFuture.kt | 180 ++++ .../kotlin/com/openai/core/ClientOptions.kt | 91 +- .../com/openai/core/ClientOptionsView.kt | 12 + .../kotlin/com/openai/core/PrepareRequest.kt | 11 +- .../com/openai/core/RequestAuthentication.kt | 16 +- .../kotlin/com/openai/core/RequestOptions.kt | 4 + .../com/openai/core/handlers/ErrorHandler.kt | 98 +- .../openai/core/http/AsyncStreamResponse.kt | 10 +- .../http/HttpRequestAttemptAuthenticator.kt | 44 + .../com/openai/core/http/HttpResponseFor.kt | 33 +- .../com/openai/core/http/LoggingHttpClient.kt | 99 +- ...eClosingHttpRequestAttemptAuthenticator.kt | 25 + .../openai/core/http/PipelineRequestBody.kt | 52 + .../openai/core/http/PipelineResponseLease.kt | 63 ++ .../openai/core/http/RetryingHttpClient.kt | 239 +---- .../http/RetryingHttpClientOrchestrator.kt | 666 ++++++++++++ .../async/ContainerServiceAsyncImpl.kt | 3 +- .../async/ResponseServiceAsyncImpl.kt | 3 +- .../async/beta/ResponseServiceAsyncImpl.kt | 3 +- .../async/containers/FileServiceAsyncImpl.kt | 3 +- .../async/realtime/CallServiceAsyncImpl.kt | 9 +- .../services/blocking/ContainerServiceImpl.kt | 3 +- .../services/blocking/ResponseServiceImpl.kt | 3 +- .../blocking/beta/ResponseServiceImpl.kt | 3 +- .../blocking/containers/FileServiceImpl.kt | 3 +- .../blocking/realtime/CallServiceImpl.kt | 9 +- .../core/CancellationPropagatingFutureTest.kt | 713 +++++++++++++ .../CancellationPropagatingRequestBodyTest.kt | 75 ++ .../com/openai/core/ClientOptionsTest.kt | 11 +- .../core/X509BlockingResponseLifecycleTest.kt | 118 +++ .../core/handlers/PipelineErrorHandlerTest.kt | 42 + ...uthenticatingRetryingHttpClientFixtures.kt | 150 +++ ...mptAuthenticatingRetryingHttpClientTest.kt | 944 ++++++++++++++++++ .../AuthenticatedRequestBodyLifecycleTest.kt | 93 ++ .../http/LoggingHttpClientCancellationTest.kt | 193 ++++ .../core/http/PipelineHttpResponseForTest.kt | 56 ++ .../core/http/PipelineResponseLeaseTest.kt | 34 + .../http/X509AsyncStreamCancellationTest.kt | 169 ++++ .../example/X509WorkloadIdentityExample.java | 77 ++ .../compatibility/OkHttpRuntimeProbe.java | 249 +++++ 49 files changed, 5362 insertions(+), 385 deletions(-) create mode 100644 .github/workflows/x509-live-smoke.yml create mode 100644 openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/X509AttemptAuthenticatorTest.kt create mode 100644 openai-java-core/src/main/kotlin/com/openai/core/CancellationPropagatingFuture.kt create mode 100644 openai-java-core/src/main/kotlin/com/openai/core/ClientOptionsView.kt create mode 100644 openai-java-core/src/main/kotlin/com/openai/core/http/HttpRequestAttemptAuthenticator.kt create mode 100644 openai-java-core/src/main/kotlin/com/openai/core/http/PhantomReachableClosingHttpRequestAttemptAuthenticator.kt create mode 100644 openai-java-core/src/main/kotlin/com/openai/core/http/PipelineRequestBody.kt create mode 100644 openai-java-core/src/main/kotlin/com/openai/core/http/PipelineResponseLease.kt create mode 100644 openai-java-core/src/main/kotlin/com/openai/core/http/RetryingHttpClientOrchestrator.kt create mode 100644 openai-java-core/src/test/kotlin/com/openai/core/CancellationPropagatingFutureTest.kt create mode 100644 openai-java-core/src/test/kotlin/com/openai/core/CancellationPropagatingRequestBodyTest.kt create mode 100644 openai-java-core/src/test/kotlin/com/openai/core/X509BlockingResponseLifecycleTest.kt create mode 100644 openai-java-core/src/test/kotlin/com/openai/core/handlers/PipelineErrorHandlerTest.kt create mode 100644 openai-java-core/src/test/kotlin/com/openai/core/http/AttemptAuthenticatingRetryingHttpClientFixtures.kt create mode 100644 openai-java-core/src/test/kotlin/com/openai/core/http/AttemptAuthenticatingRetryingHttpClientTest.kt create mode 100644 openai-java-core/src/test/kotlin/com/openai/core/http/AuthenticatedRequestBodyLifecycleTest.kt create mode 100644 openai-java-core/src/test/kotlin/com/openai/core/http/LoggingHttpClientCancellationTest.kt create mode 100644 openai-java-core/src/test/kotlin/com/openai/core/http/PipelineHttpResponseForTest.kt create mode 100644 openai-java-core/src/test/kotlin/com/openai/core/http/PipelineResponseLeaseTest.kt create mode 100644 openai-java-core/src/test/kotlin/com/openai/core/http/X509AsyncStreamCancellationTest.kt create mode 100644 openai-java-example/src/main/java/com/openai/example/X509WorkloadIdentityExample.java diff --git a/.github/workflows/x509-live-smoke.yml b/.github/workflows/x509-live-smoke.yml new file mode 100644 index 000000000..ce92a9e05 --- /dev/null +++ b/.github/workflows/x509-live-smoke.yml @@ -0,0 +1,69 @@ +name: X.509 live smoke + +on: + workflow_dispatch: + inputs: + run_x509: + description: Confirm the protected production X.509 smoke test + required: true + default: false + type: boolean + +permissions: {} + +concurrency: + group: x509-live-smoke + cancel-in-progress: false + +jobs: + smoke: + if: >- + github.repository == 'openai/openai-java' && + github.ref == 'refs/heads/main' && + inputs.run_x509 + runs-on: ubuntu-24.04 + timeout-minutes: 10 + # Configure this environment with an independent SDK-team required reviewer and no bypass. + environment: x509-live-smoke + permissions: + contents: read + + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v6 + with: + persist-credentials: false + + - name: Set up Java + uses: actions/setup-java@b6effb05e454b25005698d916606bdc6ffcbf961 # v5 + with: + distribution: temurin + java-version: 21 + + - name: Run certificate-only production smoke + shell: bash + env: + OPENAI_X509_KEYSTORE_P12_BASE64: ${{ secrets.OPENAI_X509_KEYSTORE_P12_BASE64 }} + OPENAI_X509_KEYSTORE_PASSWORD: ${{ secrets.OPENAI_X509_KEYSTORE_PASSWORD }} + OPENAI_X509_CERTIFICATE_ALIAS: ${{ secrets.OPENAI_X509_CERTIFICATE_ALIAS }} + OPENAI_X509_IDENTITY_PROVIDER_ID: ${{ secrets.OPENAI_X509_IDENTITY_PROVIDER_ID }} + OPENAI_X509_SERVICE_ACCOUNT_ID: ${{ secrets.OPENAI_X509_SERVICE_ACCOUNT_ID }} + run: | + set -euo pipefail + umask 077 + for name in \ + OPENAI_X509_KEYSTORE_P12_BASE64 \ + OPENAI_X509_KEYSTORE_PASSWORD \ + OPENAI_X509_CERTIFICATE_ALIAS \ + OPENAI_X509_IDENTITY_PROVIDER_ID \ + OPENAI_X509_SERVICE_ACCOUNT_ID; do + if [[ -z "${!name:-}" ]]; then + echo "Missing required X.509 smoke-test configuration: $name" >&2 + exit 1 + fi + done + x509_tmp_dir="$(mktemp -d "$RUNNER_TEMP/openai-x509-live.XXXXXX")" + trap 'rm -rf -- "$x509_tmp_dir"' EXIT + export OPENAI_X509_KEYSTORE_PATH="$x509_tmp_dir/identity.p12" + printf '%s' "$OPENAI_X509_KEYSTORE_P12_BASE64" | base64 --decode > "$OPENAI_X509_KEYSTORE_PATH" + chmod 600 "$OPENAI_X509_KEYSTORE_PATH" + ./gradlew :openai-java-example:run -Pexample=X509WorkloadIdentity --no-daemon diff --git a/README.md b/README.md index cc4242fd8..f8866c3a2 100644 --- a/README.md +++ b/README.md @@ -203,6 +203,45 @@ OpenAIClient client = OpenAIOkHttpClient.builder() .build(); ``` +#### X.509 workload identity federation (preview) + +X.509 workload identity exchanges a client certificate for a short-lived bearer token and then +presents the same certificate to the OpenAI mTLS API. It does not use an API key: + +```java +X509WorkloadIdentity identity = X509WorkloadIdentity.builder() + .identityProviderId("your-identity-provider-id") + .serviceAccountId("your-service-account-id") + .build(); + +X509Transport transport = X509Transport.builder() + .keyManager(keyManager) + .certificateAlias("workload-certificate") + .trustManager(trustManager) + .build(); + +OpenAIClient client = OpenAIOkHttpClient.x509Builder(identity, transport).build(); +``` + +`keyManager` must contain the private key and certificate chain for the selected alias; +`trustManager` verifies OpenAI's servers. See +[`X509WorkloadIdentityExample`](openai-java-example/src/main/java/com/openai/example/X509WorkloadIdentityExample.java) +for a complete PKCS#12 example. + +This mode deliberately fixes both network destinations: token exchange goes directly to +`https://mtls.auth.openai.com/oauth/token`, and API requests go directly to +`https://mtls.api.openai.com/v1`. It rejects API keys, admin keys, `fromEnv()`, custom base URLs, +organization/project headers, proxies, redirects, and custom transports. A client owns one +generation-scoped token cache; concurrent requests share exchanges, and transient exchange and API +failures share one retry budget and total deadline. Separately, one `401` can invalidate the exact +rejected token and replay a repeatable request once. + +The bearer token is not cryptographically certificate-bound unless the service includes and +enforces a confirmation (`cnf`) claim. Treat the token as a credential: never log it, and do not +forward it outside the fixed mTLS client. To rotate a certificate, build a new key manager, +transport, and client, atomically direct new work to that client, then drain and close the old +client. Mutating a key store behind a live client is unsupported. + #### Kubernetes service account token provider ```java diff --git a/buildSrc/src/main/kotlin/com/openai/gradle/CoreCompilationShards.kt b/buildSrc/src/main/kotlin/com/openai/gradle/CoreCompilationShards.kt index 9edc6ac12..b3dc7bf96 100644 --- a/buildSrc/src/main/kotlin/com/openai/gradle/CoreCompilationShards.kt +++ b/buildSrc/src/main/kotlin/com/openai/gradle/CoreCompilationShards.kt @@ -109,15 +109,19 @@ object CoreCompilationShards { private val clientBaseSources = setOf( + "com/openai/core/CancellationPropagatingFuture.kt", "com/openai/core/ClientOptions.kt", + "com/openai/core/ClientOptionsView.kt", "com/openai/core/PrepareRequest.kt", "com/openai/core/Properties.kt", "com/openai/core/RequestOptions.kt", "com/openai/core/http/AuthenticatingHttpClient.kt", "com/openai/core/http/HttpClient.kt", "com/openai/core/http/LoggingHttpClient.kt", + "com/openai/core/http/PipelineRequestBody.kt", "com/openai/core/http/PhantomReachableClosingHttpClient.kt", "com/openai/core/http/RetryingHttpClient.kt", + "com/openai/core/http/RetryingHttpClientOrchestrator.kt", "com/openai/core/http/WorkloadIdentityHttpClient.kt", ) diff --git a/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/X509ClientIntegration.kt b/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/X509ClientIntegration.kt index 0bf55da42..1e2c81e70 100644 --- a/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/X509ClientIntegration.kt +++ b/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/X509ClientIntegration.kt @@ -3,12 +3,23 @@ package com.openai.client.okhttp import com.openai.auth.X509WorkloadIdentity import com.openai.core.ClientOptions import com.openai.core.Timeout +import com.openai.core.http.AuthenticatedHttpRequest import com.openai.core.http.HttpRequest -import com.openai.core.http.HttpRequestAuthenticator +import com.openai.core.http.HttpRequestAttemptAuthenticator +import com.openai.errors.OpenAIIoException +import com.openai.errors.OpenAIRetryableException +import com.openai.errors.UnexpectedStatusCodeException +import java.io.IOException +import java.time.Duration import java.util.Locale -import java.util.concurrent.CancellationException import java.util.concurrent.CompletableFuture +import java.util.concurrent.CompletionException +import java.util.concurrent.ExecutionException +import java.util.concurrent.RejectedExecutionException +import java.util.concurrent.ScheduledThreadPoolExecutor +import java.util.concurrent.TimeUnit import java.util.concurrent.atomic.AtomicBoolean +import java.util.concurrent.atomic.AtomicLong internal const val X509_API_BASE_URL = "https://mtls.api.openai.com/v1" @@ -37,7 +48,7 @@ private constructor( clientOptions .fixedBearerTransport( transport.apiClient, - X509RequestAuthenticator(identity, transport.exchangeClient), + X509AttemptAuthenticator(identity, transport.exchangeClient), ) .build() } catch (error: Throwable) { @@ -53,59 +64,247 @@ private constructor( } } -/** Owns the exchange client and adds a freshly exchanged bearer to API requests. */ -private class X509RequestAuthenticator( - identity: X509WorkloadIdentity, - private val exchangeClient: OkHttpClient, -) : HttpRequestAuthenticator { - private val tokenExchange = X509TokenExchange(identity, exchangeClient) +/** Owns the exchange client and installs one exact, generation-scoped bearer per API attempt. */ +private class X509AttemptAuthenticator( + private val exchange: () -> CompletableFuture, + private val closeExchange: () -> Unit, + private val nanoTime: () -> Long, + private val beforeTokenPublication: () -> Unit, + private val beforeRefreshCleared: () -> Unit, + private val beforeWaiterTimeoutSchedule: () -> Unit, + private val scheduler: ScheduledThreadPoolExecutor, +) : HttpRequestAttemptAuthenticator { + constructor( + identity: X509WorkloadIdentity, + exchangeClient: OkHttpClient, + ) : this( + X509TokenExchange(identity, exchangeClient)::executeAsync, + exchangeClient::close, + System::nanoTime, + {}, + {}, + {}, + tokenWaitScheduler(), + ) + + private val serial = AtomicLong() + private val lock = Any() + private var cached: CachedToken? = null + private var refresh: Refresh? = null + private var invalidationEpoch = 0L private val closed = AtomicBoolean() - override fun authenticate(request: HttpRequest): HttpRequest { + override fun authenticate(request: HttpRequest, timeout: Duration?): AuthenticatedHttpRequest { validateRequest(request) - return authenticated(request, tokenExchange.execute()) + val waiter = token(timeout) + val token = + try { + waiter.get() + } catch (error: InterruptedException) { + waiter.cancel(true) + Thread.currentThread().interrupt() + throw OpenAIIoException("Interrupted while obtaining an X.509 access token", error) + } catch (error: ExecutionException) { + throw unchecked(error.cause ?: error) + } + return authenticated(request, token) } - override fun authenticateAsync(request: HttpRequest): CompletableFuture { + override fun authenticateAsync( + request: HttpRequest, + timeout: Duration?, + ): CompletableFuture { try { validateRequest(request) } catch (error: Throwable) { - return CompletableFuture().also { it.completeExceptionally(error) } + return CompletableFuture().also { + it.completeExceptionally(error) + } + } + val token = token(timeout) + val result = CompletableFuture() + token.whenComplete { value, error -> + if (error == null) result.complete(authenticated(request, value)) + else result.completeExceptionally(unwrap(error)) } + result.whenComplete { _, _ -> if (result.isCancelled) token.cancel(true) } + return result + } - val exchangeFuture = tokenExchange.executeAsync() - val result = CompletableFuture() - exchangeFuture.whenComplete { token, error -> - if (error != null) { - result.completeExceptionally(error) - } else if (token == null) { - result.completeExceptionally( - IllegalStateException("X.509 token exchange completed without a token") - ) - } else { + private fun token(timeout: Duration?): CompletableFuture { + val now = nanoTime() + val state = + synchronized(lock) { + check(!closed.get()) { "X.509 authenticator is closed" } + cached + ?.takeIf { it.isBeforeRefresh(now) && !it.isExpired(now) } + ?.let { + return CompletableFuture.completedFuture(it) + } + (refresh?.takeUnless { it.result.isDone } ?: startRefresh()).also { it.waiters++ } + } + val waiter = CompletableFuture() + val detached = AtomicBoolean() + fun detach() { + if (!detached.compareAndSet(false, true)) return + val cancel = + synchronized(lock) { + state.waiters-- + if (state.waiters == 0 && !state.result.isDone && refresh === state) { + refresh = null + true + } else false + } + if (cancel) { + state.raw.cancel(true) + state.result.cancel(true) + } + } + state.result.whenComplete { value, error -> + if (error == null) waiter.complete(value) + else waiter.completeExceptionally(unwrap(error)) + } + timeout?.let { + if (it.isZero) { + waiter.completeExceptionally(OpenAIIoException("X.509 request deadline exceeded")) + detach() + return waiter + } + val timeoutTask = try { - result.complete(authenticated(request, token)) - } catch (authenticationError: Throwable) { - result.completeExceptionally(authenticationError) + beforeWaiterTimeoutSchedule() + scheduler.schedule( + { + waiter.completeExceptionally( + OpenAIIoException("X.509 request deadline exceeded") + ) + }, + it.toNanos(), + TimeUnit.NANOSECONDS, + ) + } catch (error: RejectedExecutionException) { + waiter.completeExceptionally( + if (closed.get()) OpenAIIoException("HTTP client is closed", error) + else error + ) + detach() + return waiter + } + waiter.whenComplete { _, _ -> timeoutTask.cancel(false) } + } + waiter.whenComplete { _, _ -> detach() } + return waiter + } + + private fun startRefresh(): Refresh { + val exchangeStarted = nanoTime() + val raw = exchange() + val result = CompletableFuture() + val state = Refresh(raw, result, cached, invalidationEpoch) + refresh = state + raw.whenComplete { exchanged, rawError -> + var value: CachedToken? = null + var error = rawError?.let(::unwrap) + if (error == null) { + val lifetime = exchanged.expiresIn.toNanos() + val elapsed = elapsedSince(exchangeStarted, nanoTime()) + if (elapsed < 0 || elapsed >= lifetime) { + error = OpenAIRetryableException("X.509 access token expired during exchange") + } else { + value = + CachedToken( + exchanged.value, + serial.incrementAndGet(), + exchangeStarted, + fourFifths(lifetime), + lifetime, + ) } } + if (error == null) beforeTokenPublication() + synchronized(lock) { + if (refresh !== state) return@whenComplete + if (closed.get()) { + value = null + error = OpenAIIoException("HTTP client is closed") + } else if (value != null) { + cached = value + } else if ( + state.fallback != null && + cached === state.fallback && + invalidationEpoch == state.invalidationEpoch && + !state.fallback.isExpired(nanoTime()) && + isTransient(error) + ) { + state.fallback.deferRefresh(nanoTime(), REFRESH_FAILURE_COOLDOWN.toNanos()) + value = state.fallback + error = null + } + } + if (error == null) result.complete(requireNotNull(value)) + else result.completeExceptionally(error) + beforeRefreshCleared() + synchronized(lock) { if (refresh === state) refresh = null } } - result.whenComplete { _, error -> - if (error is CancellationException) { - exchangeFuture.cancel(true) + return state + } + + private fun authenticated(request: HttpRequest, token: CachedToken): AuthenticatedHttpRequest { + val authenticated = + request.toBuilder().replaceHeaders("Authorization", "Bearer ${token.value}").build() + return AuthenticatedHttpRequest.create(authenticated) { + synchronized(lock) { + if (cached?.serial == token.serial) { + cached = null + invalidationEpoch++ + } } } - return result } override fun close() { - if (closed.compareAndSet(false, true)) { - exchangeClient.close() - } + val active = + synchronized(lock) { + if (!closed.compareAndSet(false, true)) return + val value = refresh + invalidationEpoch++ + cached = null + refresh = null + value?.result?.completeExceptionally(OpenAIIoException("HTTP client is closed")) + value + } + active?.raw?.cancel(true) + scheduler.shutdownNow() + closeExchange() } - private fun authenticated(request: HttpRequest, token: X509AccessToken): HttpRequest = - request.toBuilder().replaceHeaders("Authorization", "Bearer ${token.value}").build() + private class Refresh( + val raw: CompletableFuture, + val result: CompletableFuture, + val fallback: CachedToken?, + val invalidationEpoch: Long, + var waiters: Int = 0, + ) + + private class CachedToken( + val value: String, + val serial: Long, + private val issuedAt: Long, + @Volatile private var refreshAfter: Long, + private val expiresAfter: Long, + ) { + fun isBeforeRefresh(now: Long): Boolean = elapsedSince(issuedAt, now) < refreshAfter + + fun isExpired(now: Long): Boolean { + val elapsed = elapsedSince(issuedAt, now) + return elapsed < 0 || elapsed >= expiresAfter + } + + fun deferRefresh(now: Long, cooldown: Long) { + val elapsed = elapsedSince(issuedAt, now) + refreshAfter = minOf(saturatedAdd(elapsed, cooldown), expiresAfter) + } + } private fun validateRequest(request: HttpRequest) { require(request.baseUrl == X509_API_BASE_URL) { @@ -120,6 +319,41 @@ private class X509RequestAuthenticator( } private companion object { + val REFRESH_FAILURE_COOLDOWN: Duration = Duration.ofSeconds(1) + + fun elapsedSince(start: Long, now: Long): Long = now - start + + fun fourFifths(value: Long): Long = (value / 5) * 4 + ((value % 5) * 4) / 5 + + fun saturatedAdd(left: Long, right: Long): Long = + if (left > Long.MAX_VALUE - right) Long.MAX_VALUE else left + right + + fun unwrap(error: Throwable?): Throwable = + when (error) { + is CompletionException, + is ExecutionException -> error.cause ?: error + null -> IllegalStateException("X.509 token exchange failed without a cause") + else -> error + } + + fun unchecked(error: Throwable): RuntimeException = + unwrap(error).let { if (it is RuntimeException) it else OpenAIIoException(cause = it) } + + fun isTransient(error: Throwable?): Boolean = + when (val cause = unwrap(error)) { + is IOException, + is OpenAIIoException, + is OpenAIRetryableException -> true + is UnexpectedStatusCodeException -> + when (cause.headers().values("X-Should-Retry").firstOrNull()) { + "true" -> true + "false" -> false + else -> + cause.statusCode() in setOf(408, 409, 429) || cause.statusCode() >= 500 + } + else -> false + } + val FORBIDDEN_HEADERS = setOf( "authorization", @@ -134,3 +368,31 @@ private class X509RequestAuthenticator( ) } } + +@JvmSynthetic +internal fun x509AttemptAuthenticatorForTest( + closeExchange: () -> Unit = {}, + nanoTime: () -> Long = System::nanoTime, + beforeTokenPublication: () -> Unit = {}, + beforeRefreshCleared: () -> Unit = {}, + beforeWaiterTimeoutSchedule: () -> Unit = {}, + schedulerObserver: (ScheduledThreadPoolExecutor) -> Unit = {}, + exchange: () -> CompletableFuture, +): HttpRequestAttemptAuthenticator { + val scheduler = tokenWaitScheduler().also(schedulerObserver) + return X509AttemptAuthenticator( + exchange, + closeExchange, + nanoTime, + beforeTokenPublication, + beforeRefreshCleared, + beforeWaiterTimeoutSchedule, + scheduler, + ) +} + +private fun tokenWaitScheduler(): ScheduledThreadPoolExecutor = + ScheduledThreadPoolExecutor(1) { task -> + Thread(task, "openai-x509-token-wait").apply { isDaemon = true } + } + .apply { removeOnCancelPolicy = true } diff --git a/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/X509TokenExchange.kt b/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/X509TokenExchange.kt index e6f73abda..53b5680d8 100644 --- a/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/X509TokenExchange.kt +++ b/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/X509TokenExchange.kt @@ -4,6 +4,7 @@ import com.fasterxml.jackson.core.JsonProcessingException import com.fasterxml.jackson.databind.DeserializationFeature import com.fasterxml.jackson.databind.JsonNode import com.openai.auth.X509WorkloadIdentity +import com.openai.core.RequestOptions import com.openai.core.http.Headers import com.openai.core.http.HttpClient import com.openai.core.http.HttpMethod @@ -33,11 +34,13 @@ internal class X509TokenExchange( private val responseReader = jsonMapper.reader().with(DeserializationFeature.FAIL_ON_TRAILING_TOKENS) - fun execute(): X509AccessToken = - httpClient.execute(request()).use { response -> parse(response) } + fun execute(requestOptions: RequestOptions = RequestOptions.none()): X509AccessToken = + httpClient.execute(request(), requestOptions).use { response -> parse(response) } - fun executeAsync(): CompletableFuture { - val responseFuture = httpClient.executeAsync(request()) + fun executeAsync( + requestOptions: RequestOptions = RequestOptions.none() + ): CompletableFuture { + val responseFuture = httpClient.executeAsync(request(), requestOptions) val result = CompletableFuture() val activeResponse = AtomicReference() responseFuture.whenCompleteAsync { response, error -> diff --git a/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/OpenAIOkHttpClientX509Test.kt b/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/OpenAIOkHttpClientX509Test.kt index 7ee728dd6..87efc234c 100644 --- a/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/OpenAIOkHttpClientX509Test.kt +++ b/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/OpenAIOkHttpClientX509Test.kt @@ -29,7 +29,7 @@ internal class OpenAIOkHttpClientX509Test { .syncBuilder() .putHeader("X-Application-Test", "api-only") .build() - .withOptions { options -> options.timeout(Duration.ofSeconds(5)) } + .withOptions { options -> options.timeout(Duration.ofSeconds(30)) } try { assertThat(client.files().list().data()).isEmpty() @@ -61,6 +61,194 @@ internal class OpenAIOkHttpClientX509Test { } } + @Test + fun syncClientAsyncViewUsesAsyncUserAgentAndSharedLifecycle() { + Fixture().use { fixture -> + fixture.enqueueSuccess() + val client = fixture.syncBuilder().build().async() + + try { + assertThat(client.files().list().get(10, TimeUnit.SECONDS).data()).isEmpty() + } finally { + client.close() + } + + assertSuccessfulWireContract(fixture, expectedUserAgent = "OpenAIClientAsyncImpl/Java") + } + } + + @Test + fun asyncClientSyncViewUsesSyncUserAgentAndSharedLifecycle() { + Fixture().use { fixture -> + fixture.enqueueSuccess() + val client = fixture.asyncBuilder().build().sync() + + try { + assertThat(client.files().list().data()).isEmpty() + } finally { + client.close() + } + + assertSuccessfulWireContract(fixture, expectedUserAgent = "OpenAIClientImpl/Java") + } + } + + @Test + fun syncBuilderCanBuildIndependentClients() { + Fixture().use { fixture -> + fixture.enqueueSuccess() + fixture.enqueueSuccess() + val builder = fixture.syncBuilder() + val first = builder.build() + val second = builder.build() + + assertThat(first).isNotSameAs(second) + try { + assertThat(first.files().list().data()).isEmpty() + } finally { + first.close() + } + try { + assertThat(second.files().list().data()).isEmpty() + } finally { + second.close() + } + } + } + + @Test + fun asyncBuilderCanBuildIndependentClients() { + Fixture().use { fixture -> + fixture.enqueueSuccess() + fixture.enqueueSuccess() + val builder = fixture.asyncBuilder() + val first = builder.build() + val second = builder.build() + + assertThat(first).isNotSameAs(second) + try { + assertThat(first.files().list().get(10, TimeUnit.SECONDS).data()).isEmpty() + } finally { + first.close() + } + try { + assertThat(second.files().list().get(10, TimeUnit.SECONDS).data()).isEmpty() + } finally { + second.close() + } + } + } + + @Test + fun reusableBuilderDoesNotPersistDefaultUserAgent() { + Fixture().use { fixture -> + fixture.enqueueSuccess() + fixture.enqueueSuccess() + val builder = fixture.syncBuilder() + val first = builder.build() + val second = builder.putHeader("User-Agent", "caller-agent").build() + + try { + assertThat(first.files().list().data()).isEmpty() + } finally { + first.close() + } + assertSuccessfulWireContract(fixture, expectedUserAgent = "OpenAIClientImpl/Java") + + try { + assertThat(second.files().list().data()).isEmpty() + } finally { + second.close() + } + assertSuccessfulWireContract(fixture, expectedUserAgent = "caller-agent") + } + } + + @Test + fun reusesCachedTokenAcrossSequentialApiRequests() { + Fixture().use { fixture -> + fixture.enqueueExchange(ACCESS_TOKEN, closeConnection = false) + fixture.enqueueApiSuccess(closeConnection = false) + fixture.apiPeer.server.enqueue(fixture.apiSuccess()) + val client = fixture.syncBuilder().build() + + try { + assertThat(client.files().list().data()).isEmpty() + assertThat(client.files().list().data()).isEmpty() + } finally { + client.close() + } + + assertThat(fixture.authPeer.server.requestCount).isEqualTo(2) + fixture.authPeer.takeRequest() + fixture.authPeer.takeRequest() + fixture.apiPeer.takeRequest() + val apiRequests = listOf(fixture.apiPeer.takeRequest(), fixture.apiPeer.takeRequest()) + assertThat(apiRequests.map { it.getHeader("Authorization") }) + .containsExactly("Bearer $ACCESS_TOKEN", "Bearer $ACCESS_TOKEN") + } + } + + @Test + fun unauthorizedResponseInvalidatesExactTokenAndReplaysOnceWithRetriesDisabled() { + Fixture().use { fixture -> + fixture.enqueueExchange("tokenone", closeConnection = false) + fixture.apiPeer.enqueue(MockResponse().setResponseCode(401)) + fixture.authPeer.server.enqueue(fixture.exchangeResponse("tokentwo")) + fixture.apiPeer.server.enqueue(fixture.apiSuccess()) + val client = fixture.asyncBuilder().maxRetries(0).build() + + try { + assertThat(client.files().list().get(10, TimeUnit.SECONDS).data()).isEmpty() + } finally { + client.close() + } + + fixture.authPeer.takeRequest() + val exchangeRequests = List(2) { fixture.authPeer.takeRequest() } + fixture.apiPeer.takeRequest() + val apiRequests = List(2) { fixture.apiPeer.takeRequest() } + assertThat(exchangeRequests).allMatch { it.path == "/oauth/token" } + assertThat(apiRequests.map { it.getHeader("Authorization") }) + .containsExactly("Bearer tokenone", "Bearer tokentwo") + assertThat(apiRequests.map { it.getHeader("X-Stainless-Retry-Count") }) + .containsExactly("0", "0") + } + } + + @Test + fun transientFailuresBeforeAndAfterUnauthorizedShareRetryBudget() { + Fixture().use { fixture -> + fixture.enqueueExchange("tokenone", closeConnection = false) + fixture.apiPeer.enqueue(fixture.apiFailure(500)) + fixture.apiPeer.server.enqueue(MockResponse().setResponseCode(401)) + fixture.authPeer.server.enqueue(fixture.exchangeResponse("tokentwo")) + fixture.apiPeer.server.enqueue(fixture.apiFailure(500)) + fixture.apiPeer.server.enqueue(fixture.apiFailure(500)) + val client = fixture.syncBuilder().maxRetries(2).sleeper(NoDelaySleeper).build() + + try { + assertThatThrownBy { client.files().list() } + .isInstanceOf(RuntimeException::class.java) + } finally { + client.close() + } + + fixture.apiPeer.takeRequest() + val apiRequests = List(4) { fixture.apiPeer.takeRequest() } + assertThat(apiRequests.map { it.getHeader("Authorization") }) + .containsExactly( + "Bearer tokenone", + "Bearer tokenone", + "Bearer tokentwo", + "Bearer tokentwo", + ) + assertThat(apiRequests.map { it.getHeader("X-Stainless-Retry-Count") }) + .containsExactly("0", "1", "1", "2") + assertThat(fixture.apiPeer.server.requestCount).isEqualTo(5) + } + } + @Test fun onlyPublicJavaConstructionPathIsX509BuilderFactory() { listOf(OpenAIOkHttpClient.Builder::class.java, OpenAIOkHttpClientAsync.Builder::class.java) @@ -240,6 +428,26 @@ internal class OpenAIOkHttpClientX509Test { } } + @Test + fun cancellingPublicAsyncFutureCancelsBlockedExchange() { + Fixture().use { fixture -> + fixture.authPeer.enqueue(MockResponse().setSocketPolicy(SocketPolicy.NO_RESPONSE)) + val client = fixture.asyncBuilder().timeout(Duration.ofSeconds(30)).build() + val cancelled = client.files().list() + + assertConnectAuthority(fixture.authPeer.takeRequest(), AUTH_AUTHORITY) + assertThat(fixture.authPeer.takeRequest().path).isEqualTo("/oauth/token") + cancelled.cancel(true) + + try { + assertThat(cancelled.isCancelled).isTrue() + assertThat(fixture.apiPeer.server.requestCount).isZero() + } finally { + client.close() + } + } + } + @Test fun closingSyncClientCancelsBlockedExchangeBeforeApiDispatch() { Fixture().use { fixture -> @@ -377,6 +585,41 @@ internal class OpenAIOkHttpClientX509Test { ) } + fun enqueueExchange(token: String, closeConnection: Boolean = true) { + authPeer.enqueue( + exchangeResponse(token).apply { + if (closeConnection) setHeader("Connection", "close") + } + ) + } + + fun exchangeResponse(token: String): MockResponse = + MockResponse() + .setHeader("Content-Type", "application/json") + .setBody(tokenResponse(token)) + + fun enqueueApiSuccess(closeConnection: Boolean = true) { + val response = apiSuccess() + if (closeConnection) enqueueApi(response) + else apiPeer.enqueue(response.removeHeader("Connection")) + } + + fun enqueueApi(response: MockResponse) { + apiPeer.enqueue(response.setHeader("Connection", "close")) + } + + fun apiSuccess(): MockResponse = + MockResponse() + .setHeader("Content-Type", "application/json") + .setHeader("Connection", "close") + .setBody(FILES_RESPONSE) + + fun apiFailure(status: Int): MockResponse = + MockResponse() + .setResponseCode(status) + .setHeader("Content-Type", "application/json") + .setBody("""{"error":{"message":"test","type":"server_error"}}""") + override fun close() { apiPeer.use { authPeer.close() } } @@ -411,5 +654,16 @@ internal class OpenAIOkHttpClientX509Test { } """ .trimIndent() + + fun tokenResponse(token: String): String = TOKEN_RESPONSE.replace(ACCESS_TOKEN, token) + + private object NoDelaySleeper : com.openai.core.Sleeper { + override fun sleep(duration: Duration) {} + + override fun sleepAsync(duration: Duration) = + java.util.concurrent.CompletableFuture.completedFuture(null) + + override fun close() {} + } } } diff --git a/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/X509AttemptAuthenticatorTest.kt b/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/X509AttemptAuthenticatorTest.kt new file mode 100644 index 000000000..3b0b66692 --- /dev/null +++ b/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/X509AttemptAuthenticatorTest.kt @@ -0,0 +1,411 @@ +package com.openai.client.okhttp + +import com.openai.core.RequestOptions +import com.openai.core.http.Headers +import com.openai.core.http.HttpClient +import com.openai.core.http.HttpMethod +import com.openai.core.http.HttpRequest +import com.openai.core.http.HttpResponse +import com.openai.core.http.RetryingHttpClient +import com.openai.errors.OpenAIRetryableException +import java.io.ByteArrayInputStream +import java.time.Duration +import java.util.concurrent.CompletableFuture +import java.util.concurrent.CountDownLatch +import java.util.concurrent.ExecutionException +import java.util.concurrent.Executors +import java.util.concurrent.ScheduledThreadPoolExecutor +import java.util.concurrent.TimeUnit +import java.util.concurrent.atomic.AtomicInteger +import java.util.concurrent.atomic.AtomicLong +import org.assertj.core.api.Assertions.assertThat +import org.assertj.core.api.Assertions.assertThatThrownBy +import org.junit.jupiter.api.Test + +internal class X509AttemptAuthenticatorTest { + @Test + fun oneCancelledWaiterDoesNotCancelSharedExchange() { + val exchange = CancellationFuture() + val authenticator = x509AttemptAuthenticatorForTest { exchange } + val cancelled = authenticator.authenticateAsync(request(), Duration.ofSeconds(5)) + val surviving = authenticator.authenticateAsync(request(), Duration.ofSeconds(5)) + + cancelled.cancel(true) + exchange.complete(X509AccessToken("survivingtoken", Duration.ofMinutes(1))) + + assertThat(exchange.isCancelled).isFalse() + assertThat(authorization(surviving.get(5, TimeUnit.SECONDS))) + .isEqualTo("Bearer survivingtoken") + authenticator.close() + } + + @Test + fun lastCancelledWaiterCancelsExchangeAndCannotInstallLateToken() { + val first = CancellationFuture() + val second = CancellationFuture() + val exchanges = ArrayDeque(listOf(first, second)) + val authenticator = x509AttemptAuthenticatorForTest { exchanges.removeFirst() } + val cancelled = authenticator.authenticateAsync(request(), Duration.ofSeconds(5)) + + cancelled.cancel(true) + + assertThat(first.cancelled.await(5, TimeUnit.SECONDS)).isTrue() + assertThat(first.complete(X509AccessToken("latetoken", Duration.ofMinutes(1)))).isFalse() + val replacement = authenticator.authenticateAsync(request(), Duration.ofSeconds(5)) + second.complete(X509AccessToken("replacementtoken", Duration.ofMinutes(1))) + assertThat(authorization(replacement.get(5, TimeUnit.SECONDS))) + .isEqualTo("Bearer replacementtoken") + authenticator.close() + } + + @Test + fun timingOutOneWaiterLeavesTheSharedExchangeForOtherWaiters() { + val exchange = CancellationFuture() + val authenticator = x509AttemptAuthenticatorForTest { exchange } + val surviving = authenticator.authenticateAsync(request(), Duration.ofSeconds(5)) + val timedOut = authenticator.authenticateAsync(request(), Duration.ofMillis(20)) + + assertThatThrownBy { timedOut.get(5, TimeUnit.SECONDS) } + .isInstanceOf(ExecutionException::class.java) + .hasMessageContaining("deadline") + assertThat(exchange.isCancelled).isFalse() + exchange.complete(X509AccessToken("survivingtoken", Duration.ofMinutes(1))) + assertThat(authorization(surviving.get(5, TimeUnit.SECONDS))) + .isEqualTo("Bearer survivingtoken") + authenticator.close() + } + + @Test + fun interruptingSyncWaiterCancelsItsLastExchange() { + val exchange = CancellationFuture() + val started = CountDownLatch(1) + val authenticator = x509AttemptAuthenticatorForTest { + started.countDown() + exchange + } + val executor = Executors.newSingleThreadExecutor() + val waiter = executor.submit { authenticator.authenticate(request(), null) } + + try { + assertThat(started.await(5, TimeUnit.SECONDS)).isTrue() + waiter.cancel(true) + assertThat(exchange.cancelled.await(5, TimeUnit.SECONDS)).isTrue() + } finally { + authenticator.close() + executor.shutdownNow() + } + } + + @Test + fun failedRefreshCannotRestoreConcurrentlyRejectedFallback() { + val now = AtomicLong() + val initial = CompletableFuture() + val refresh = CompletableFuture() + val replacement = CompletableFuture() + val exchanges = ArrayDeque(listOf(initial, refresh, replacement)) + val authenticator = + x509AttemptAuthenticatorForTest(nanoTime = now::get) { exchanges.removeFirst() } + val installed = authenticator.authenticateAsync(request(), Duration.ofSeconds(5)) + initial.complete(X509AccessToken("fallbacktoken", Duration.ofMillis(500))) + val rejected = installed.get(5, TimeUnit.SECONDS) + now.set(Duration.ofMillis(425).toNanos()) + + val refreshing = authenticator.authenticateAsync(request(), Duration.ofSeconds(5)) + rejected.onUnauthorized() + refresh.completeExceptionally(OpenAIRetryableException("refresh failed")) + + assertThatThrownBy { refreshing.get(5, TimeUnit.SECONDS) } + .isInstanceOf(ExecutionException::class.java) + .hasMessageContaining("refresh failed") + val afterRejection = authenticator.authenticateAsync(request(), Duration.ofSeconds(5)) + replacement.complete(X509AccessToken("replacementtoken", Duration.ofMinutes(1))) + assertThat(authorization(afterRejection.get(5, TimeUnit.SECONDS))) + .isEqualTo("Bearer replacementtoken") + authenticator.close() + } + + @Test + fun refreshFailureCooldownCannotServeFallbackPastExpiry() { + val now = AtomicLong() + val initial = CompletableFuture() + val refresh = CompletableFuture() + val replacement = CompletableFuture() + val exchanges = ArrayDeque(listOf(initial, refresh, replacement)) + val authenticator = + x509AttemptAuthenticatorForTest(nanoTime = now::get) { exchanges.removeFirst() } + val installed = authenticator.authenticateAsync(request(), Duration.ofSeconds(5)) + initial.complete(X509AccessToken("fallbacktoken", Duration.ofMillis(500))) + installed.get(5, TimeUnit.SECONDS) + now.set(Duration.ofMillis(425).toNanos()) + + val fallback = authenticator.authenticateAsync(request(), Duration.ofSeconds(5)) + refresh.completeExceptionally(OpenAIRetryableException("refresh failed")) + assertThat(authorization(fallback.get(5, TimeUnit.SECONDS))) + .isEqualTo("Bearer fallbacktoken") + now.set(Duration.ofMillis(501).toNanos()) + + val afterExpiry = authenticator.authenticateAsync(request(), Duration.ofSeconds(5)) + assertThat(afterExpiry.isDone).isFalse() + replacement.complete(X509AccessToken("replacementtoken", Duration.ofMinutes(1))) + assertThat(authorization(afterExpiry.get(5, TimeUnit.SECONDS))) + .isEqualTo("Bearer replacementtoken") + authenticator.close() + } + + @Test + fun delayedExchangeCannotExtendTokenLifetime() { + val now = AtomicLong() + val exchange = CompletableFuture() + val authenticator = x509AttemptAuthenticatorForTest(nanoTime = now::get) { exchange } + val authentication = authenticator.authenticateAsync(request(), Duration.ofSeconds(5)) + + now.set(Duration.ofSeconds(2).toNanos()) + exchange.complete(X509AccessToken("expiredtoken", Duration.ofSeconds(1))) + + assertThatThrownBy { authentication.get(5, TimeUnit.SECONDS) } + .isInstanceOf(ExecutionException::class.java) + .hasMessageContaining("expired during exchange") + authenticator.close() + } + + @Test + fun exchangeCannotPublishTokenAfterCloseStarts() { + val exchange = CompletableFuture() + lateinit var authenticator: com.openai.core.http.HttpRequestAttemptAuthenticator + authenticator = + x509AttemptAuthenticatorForTest(beforeTokenPublication = { authenticator.close() }) { + exchange + } + val authentication = authenticator.authenticateAsync(request(), Duration.ofSeconds(5)) + + exchange.complete(X509AccessToken("latetoken", Duration.ofMinutes(1))) + + assertThatThrownBy { authentication.get(5, TimeUnit.SECONDS) } + .isInstanceOf(ExecutionException::class.java) + .hasMessageContaining("HTTP client is closed") + } + + @Test + fun closeBeforeTimeoutSchedulingDoesNotExposeSchedulerRejection() { + val exchange = CancellationFuture() + lateinit var authenticator: com.openai.core.http.HttpRequestAttemptAuthenticator + authenticator = + x509AttemptAuthenticatorForTest( + beforeWaiterTimeoutSchedule = { authenticator.close() } + ) { + exchange + } + + val authentication = authenticator.authenticateAsync(request(), Duration.ofSeconds(5)) + + assertThatThrownBy { authentication.get(5, TimeUnit.SECONDS) } + .isInstanceOf(ExecutionException::class.java) + .hasMessageContaining("HTTP client is closed") + assertThat(exchange.isCancelled).isTrue() + } + + @Test + fun cachedTokenRemainsFreshAcrossNanoTimeRollover() { + val now = AtomicLong(Long.MAX_VALUE - 100) + var exchanges = 0 + val authenticator = + x509AttemptAuthenticatorForTest(nanoTime = now::get) { + exchanges++ + CompletableFuture.completedFuture( + X509AccessToken("rollovertoken", Duration.ofNanos(500)) + ) + } + + assertThat( + authorization( + authenticator + .authenticateAsync(request(), Duration.ofSeconds(5)) + .get(5, TimeUnit.SECONDS) + ) + ) + .isEqualTo("Bearer rollovertoken") + now.addAndGet(200) + assertThat( + authorization( + authenticator + .authenticateAsync(request(), Duration.ofSeconds(5)) + .get(5, TimeUnit.SECONDS) + ) + ) + .isEqualTo("Bearer rollovertoken") + assertThat(exchanges).isEqualTo(1) + authenticator.close() + } + + @Test + fun refreshFailureCooldownAndExpiryRemainOrderedAcrossNanoTimeRollover() { + val now = AtomicLong(Long.MAX_VALUE - 100) + val refresh = CompletableFuture() + val replacement = CompletableFuture() + val exchanges = + ArrayDeque( + listOf( + CompletableFuture.completedFuture( + X509AccessToken("fallbacktoken", Duration.ofNanos(500)) + ), + refresh, + replacement, + ) + ) + val authenticator = + x509AttemptAuthenticatorForTest(nanoTime = now::get) { exchanges.removeFirst() } + authenticator.authenticateAsync(request(), Duration.ofSeconds(5)).get(5, TimeUnit.SECONDS) + now.addAndGet(425) + + val fallback = authenticator.authenticateAsync(request(), Duration.ofSeconds(5)) + refresh.completeExceptionally(OpenAIRetryableException("refresh failed")) + assertThat(authorization(fallback.get(5, TimeUnit.SECONDS))) + .isEqualTo("Bearer fallbacktoken") + now.addAndGet(50) + assertThat( + authorization( + authenticator + .authenticateAsync(request(), Duration.ofSeconds(5)) + .get(5, TimeUnit.SECONDS) + ) + ) + .isEqualTo("Bearer fallbacktoken") + now.addAndGet(26) + + val afterExpiry = authenticator.authenticateAsync(request(), Duration.ofSeconds(5)) + assertThat(afterExpiry.isDone).isFalse() + replacement.complete(X509AccessToken("replacementtoken", Duration.ofMinutes(1))) + assertThat(authorization(afterExpiry.get(5, TimeUnit.SECONDS))) + .isEqualTo("Bearer replacementtoken") + authenticator.close() + } + + @Test + fun rejectedCompletedGenerationIsNotReusedByImmediateReplay() { + listOf(false, true).forEach { async -> + val firstExchange = CompletableFuture() + val firstExchangeStarted = CountDownLatch(1) + val replacementExchangeStarted = CountDownLatch(1) + val publishedBeforeClear = CountDownLatch(1) + val allowRefreshClear = CountDownLatch(1) + val exchangeCalls = AtomicInteger() + val publicationCalls = AtomicInteger() + val authenticator = + x509AttemptAuthenticatorForTest( + beforeRefreshCleared = { + if (publicationCalls.incrementAndGet() == 1) { + publishedBeforeClear.countDown() + allowRefreshClear.await(5, TimeUnit.SECONDS) + } + } + ) { + if (exchangeCalls.getAndIncrement() == 0) { + firstExchangeStarted.countDown() + firstExchange + } else { + replacementExchangeStarted.countDown() + CompletableFuture.completedFuture( + X509AccessToken("replacementtoken", Duration.ofMinutes(1)) + ) + } + } + val transport = ImmediateUnauthorizedThenSuccessClient() + val client = + RetryingHttpClient.builder() + .httpClient(transport) + .attemptAuthenticator(authenticator) + .maxRetries(0) + .build() + val executor = Executors.newSingleThreadExecutor() + val completionExecutor = Executors.newSingleThreadExecutor() + + try { + val result = + if (async) { + client.executeAsync(request()) + } else { + CompletableFuture.supplyAsync({ client.execute(request()) }, executor) + } + assertThat(firstExchangeStarted.await(5, TimeUnit.SECONDS)).isTrue() + val completion = + completionExecutor.submit { + firstExchange.complete( + X509AccessToken("rejectedtoken", Duration.ofMinutes(1)) + ) + } + assertThat(publishedBeforeClear.await(5, TimeUnit.SECONDS)).isTrue() + assertThat(replacementExchangeStarted.await(5, TimeUnit.SECONDS)).isTrue() + allowRefreshClear.countDown() + completion.get(5, TimeUnit.SECONDS) + val response = result.get(5, TimeUnit.SECONDS) + response.close() + + assertThat(transport.authorization) + .containsExactly("Bearer rejectedtoken", "Bearer replacementtoken") + } finally { + allowRefreshClear.countDown() + client.close() + executor.shutdownNow() + completionExecutor.shutdownNow() + } + } + } + + @Test + fun completedWaitersAreRemovedFromSchedulerQueue() { + val exchange = CompletableFuture() + lateinit var scheduler: ScheduledThreadPoolExecutor + val authenticator = + x509AttemptAuthenticatorForTest(schedulerObserver = { scheduler = it }) { exchange } + val waiters = + List(100) { authenticator.authenticateAsync(request(), Duration.ofMinutes(10)) } + + exchange.complete(X509AccessToken("sharedtoken", Duration.ofMinutes(1))) + waiters.forEach { it.get(5, TimeUnit.SECONDS) } + + assertThat(scheduler.removeOnCancelPolicy).isTrue() + assertThat(scheduler.queue).isEmpty() + authenticator.close() + } + + private fun request(): HttpRequest = + HttpRequest.builder().method(HttpMethod.GET).baseUrl(X509_API_BASE_URL).build() + + private fun authorization(authenticated: com.openai.core.http.AuthenticatedHttpRequest) = + authenticated.request().headers.values("Authorization").single() + + private class CancellationFuture : CompletableFuture() { + val cancelled = CountDownLatch(1) + + override fun cancel(mayInterruptIfRunning: Boolean): Boolean = + super.cancel(mayInterruptIfRunning).also { if (it) cancelled.countDown() } + } + + private class ImmediateUnauthorizedThenSuccessClient : HttpClient { + val authorization = mutableListOf() + + override fun execute(request: HttpRequest, requestOptions: RequestOptions): HttpResponse = + next(request) + + override fun executeAsync( + request: HttpRequest, + requestOptions: RequestOptions, + ): CompletableFuture = CompletableFuture.completedFuture(next(request)) + + private fun next(request: HttpRequest): HttpResponse { + authorization += request.headers.values("Authorization").single() + val status = if (authorization.size == 1) 401 else 200 + return object : HttpResponse { + override fun statusCode(): Int = status + + override fun headers(): Headers = Headers.builder().build() + + override fun body() = ByteArrayInputStream(ByteArray(0)) + + override fun close() {} + } + } + + override fun close() {} + } +} diff --git a/openai-java-core/src/main/kotlin/com/openai/client/OpenAIClientAsyncImpl.kt b/openai-java-core/src/main/kotlin/com/openai/client/OpenAIClientAsyncImpl.kt index 77c7ddeeb..0bf0a1735 100644 --- a/openai-java-core/src/main/kotlin/com/openai/client/OpenAIClientAsyncImpl.kt +++ b/openai-java-core/src/main/kotlin/com/openai/client/OpenAIClientAsyncImpl.kt @@ -4,6 +4,7 @@ package com.openai.client import com.openai.core.ClientOptions import com.openai.core.getPackageVersion +import com.openai.core.withDefaultUserAgent import com.openai.services.async.AdminServiceAsync import com.openai.services.async.AdminServiceAsyncImpl import com.openai.services.async.AudioServiceAsync @@ -59,10 +60,9 @@ class OpenAIClientAsyncImpl(private val clientOptions: ClientOptions) : OpenAICl private val clientOptionsWithUserAgent = if (clientOptions.headers.names().contains("User-Agent")) clientOptions else - clientOptions - .toBuilder() - .putHeader("User-Agent", "${javaClass.simpleName}/Java ${getPackageVersion()}") - .build() + clientOptions.withDefaultUserAgent( + "${javaClass.simpleName}/Java ${getPackageVersion()}" + ) // Pass the original clientOptions so that this client sets its own User-Agent. private val sync: OpenAIClient by lazy { OpenAIClientImpl(clientOptions) } diff --git a/openai-java-core/src/main/kotlin/com/openai/client/OpenAIClientImpl.kt b/openai-java-core/src/main/kotlin/com/openai/client/OpenAIClientImpl.kt index d5e487234..0f8e88cee 100644 --- a/openai-java-core/src/main/kotlin/com/openai/client/OpenAIClientImpl.kt +++ b/openai-java-core/src/main/kotlin/com/openai/client/OpenAIClientImpl.kt @@ -4,6 +4,7 @@ package com.openai.client import com.openai.core.ClientOptions import com.openai.core.getPackageVersion +import com.openai.core.withDefaultUserAgent import com.openai.services.blocking.AdminService import com.openai.services.blocking.AdminServiceImpl import com.openai.services.blocking.AudioService @@ -59,10 +60,9 @@ class OpenAIClientImpl(private val clientOptions: ClientOptions) : OpenAIClient private val clientOptionsWithUserAgent = if (clientOptions.headers.names().contains("User-Agent")) clientOptions else - clientOptions - .toBuilder() - .putHeader("User-Agent", "${javaClass.simpleName}/Java ${getPackageVersion()}") - .build() + clientOptions.withDefaultUserAgent( + "${javaClass.simpleName}/Java ${getPackageVersion()}" + ) // Pass the original clientOptions so that this client sets its own User-Agent. private val async: OpenAIClientAsync by lazy { OpenAIClientAsyncImpl(clientOptions) } diff --git a/openai-java-core/src/main/kotlin/com/openai/core/CancellationPropagatingFuture.kt b/openai-java-core/src/main/kotlin/com/openai/core/CancellationPropagatingFuture.kt new file mode 100644 index 000000000..18051d8ca --- /dev/null +++ b/openai-java-core/src/main/kotlin/com/openai/core/CancellationPropagatingFuture.kt @@ -0,0 +1,180 @@ +package com.openai.core + +import com.openai.core.http.HttpRequest +import com.openai.core.http.PipelineOwnedResource +import com.openai.core.http.PropagatesCancellationToUpstream +import com.openai.core.http.discardPipelineBody +import java.util.concurrent.CompletableFuture +import java.util.concurrent.CompletionStage +import java.util.concurrent.Executor +import java.util.concurrent.Future +import java.util.concurrent.atomic.AtomicReference +import java.util.function.BiConsumer +import java.util.function.Consumer +import java.util.function.Function + +internal class CancellationPropagatingFuture +private constructor( + private val asyncExecutor: Executor? = null, + private val cancelUpstream: () -> Unit, +) : CompletableFuture(), PropagatesCancellationToUpstream { + override fun cancel(mayInterruptIfRunning: Boolean): Boolean { + val cancelled = super.cancel(mayInterruptIfRunning) + if (cancelled) cancelUpstream() + return cancelled + } + + override fun thenApply(function: Function): CompletableFuture { + val input = InFlightCloseable() + val result = + CancellationPropagatingFuture { + cancel(true) + input.close() + } + whenComplete { value, error -> + if (result.isDone) { + discard(value) + return@whenComplete + } + if (error != null) result.completeExceptionally(error) + else { + input.acquire(value) + if (result.isDone) { + input.close() + return@whenComplete + } + try { + val mapped = function.apply(value) + if (!result.complete(mapped)) discard(mapped) + } catch (failure: Throwable) { + result.completeExceptionally(failure) + } finally { + input.release(value) + } + } + } + return result + } + + override fun thenAccept(action: Consumer): CompletableFuture { + val input = InFlightCloseable() + val result = + CancellationPropagatingFuture { + cancel(true) + input.close() + } + whenComplete { value, error -> + if (result.isDone) { + discard(value) + return@whenComplete + } + if (error != null) result.completeExceptionally(error) + else { + input.acquire(value) + if (result.isDone) { + input.close() + return@whenComplete + } + try { + action.accept(value) + result.complete(null) + } catch (failure: Throwable) { + result.completeExceptionally(failure) + } finally { + input.release(value) + } + } + } + return result + } + + override fun thenComposeAsync( + function: Function> + ): CompletableFuture { + val active = AtomicReference>(this) + val input = InFlightCloseable() + val result = + CancellationPropagatingFuture { + active.get().cancel(true) + input.close() + } + val completion = + BiConsumer { value, error -> + if (result.isDone) { + if (error == null) discard(value) + return@BiConsumer + } + if (error != null) { + result.completeExceptionally(error) + return@BiConsumer + } + input.acquire(value) + if (result.isDone) { + input.close() + return@BiConsumer + } + val future = + try { + function.apply(value).toCompletableFuture() + } catch (failure: Throwable) { + discardPreparedRequest(value) + result.completeExceptionally(failure) + return@BiConsumer + } finally { + input.release(value) + } + active.set(future) + if (result.isDone) { + if (!future.cancel(true)) { + future.whenComplete { composedValue, composedError -> + if (composedError == null) discard(composedValue) + } + } + return@BiConsumer + } + future.whenComplete { composedValue, composedError -> + if (composedError == null) { + if (!result.complete(composedValue)) discard(composedValue) + } else result.completeExceptionally(composedError) + } + } + if (asyncExecutor == null) whenCompleteAsync(completion) + else whenCompleteAsync(completion, asyncExecutor) + return result + } + + companion object { + private class InFlightCloseable { + private val value = AtomicReference() + + fun acquire(candidate: Any?) { + if (candidate is HttpRequest || candidate is PipelineOwnedResource) { + value.set(candidate) + } + } + + fun release(candidate: Any?) { + value.compareAndSet(candidate, null) + } + + fun close() = discard(value.getAndSet(null)) + } + + private fun discard(value: Any?) { + try { + if (value is HttpRequest) value.discardPipelineBody() + else (value as? AutoCloseable)?.close() + } catch (_: Exception) {} + } + + private fun discardPreparedRequest(value: Any?) { + if (value is HttpRequest) value.discardPipelineBody() + } + + fun completed( + value: T, + asyncExecutor: Executor? = null, + ): CancellationPropagatingFuture = + CancellationPropagatingFuture(asyncExecutor) {}.apply { complete(value) } + } +} diff --git a/openai-java-core/src/main/kotlin/com/openai/core/ClientOptions.kt b/openai-java-core/src/main/kotlin/com/openai/core/ClientOptions.kt index 85c70d186..169ce7e7b 100644 --- a/openai-java-core/src/main/kotlin/com/openai/core/ClientOptions.kt +++ b/openai-java-core/src/main/kotlin/com/openai/core/ClientOptions.kt @@ -11,6 +11,7 @@ import com.openai.core.http.AsyncStreamResponse import com.openai.core.http.AuthenticatingHttpClient import com.openai.core.http.Headers import com.openai.core.http.HttpClient +import com.openai.core.http.HttpRequestAttemptAuthenticator import com.openai.core.http.HttpRequestAuthenticator import com.openai.core.http.LoggingHttpClient import com.openai.core.http.PhantomReachableClosingHttpClient @@ -154,6 +155,10 @@ private constructor( */ fun baseUrl(): String = baseUrl ?: PRODUCTION_URL + @JvmSynthetic + internal fun propagatesAsyncCancellation(): Boolean = + requestAuthentication.attemptAuthenticator != null + fun apiKey(): Optional = Optional.ofNullable(apiKey) fun adminApiKey(): Optional = Optional.ofNullable(adminApiKey) @@ -220,6 +225,7 @@ private constructor( private var project: String? = null private var webhookSecret: String? = null private var workloadIdentity: WorkloadIdentity? = null + private var sharedHttpPipeline: HttpClient? = null @JvmSynthetic internal fun from(clientOptions: ClientOptions) = apply { @@ -255,6 +261,11 @@ private constructor( webhookSecret = clientOptions.webhookSecret } + @JvmSynthetic + internal fun shareHttpPipeline(httpClient: HttpClient) = apply { + sharedHttpPipeline = httpClient + } + /** * The HTTP client to use in the SDK. * @@ -311,14 +322,20 @@ private constructor( @JvmSynthetic fun fixedBearerTransport( httpClient: HttpClient, - httpRequestAuthenticator: HttpRequestAuthenticator, + httpRequestAuthenticator: HttpRequestAttemptAuthenticator, ) = apply { - val reserved = requestAuthentication as? RequestAuthentication.FixedBearerReserved - checkNotNull(reserved) { "Fixed bearer authentication must be set first" } + val fixedBaseUrl = + when (val authentication = requestAuthentication) { + is RequestAuthentication.FixedBearerReserved -> + authentication.fixedBearerBaseUrl + is RequestAuthentication.FixedBearerInstalled -> + authentication.fixedBearerBaseUrl + else -> error("Fixed bearer authentication must be set first") + } this.httpClient = PhantomReachableClosingHttpClient(httpClient) requestAuthentication = RequestAuthentication.FixedBearerInstalled.create( - reserved.fixedBearerBaseUrl, + fixedBaseUrl, httpRequestAuthenticator, ) } @@ -834,36 +851,44 @@ private constructor( val effectiveWorkloadIdentityAuth = (credential as? WorkloadIdentityCredential)?.getAuth() - val loggingDelegate = - if (requestAuthentication.authenticator != null) httpClient - else - WorkloadIdentityHttpClient( - delegate = httpClient, - workloadIdentityAuth = effectiveWorkloadIdentityAuth, - ) - - val loggingHttpClient = - LoggingHttpClient.builder() - .httpClient(loggingDelegate) - .clock(clock) - .level(logLevel) - .build() - - val perAttemptHttpClient = - requestAuthentication.authenticator?.let { authenticator -> - AuthenticatingHttpClient( - delegate = loggingHttpClient, - authenticator = authenticator, - ) - } ?: loggingHttpClient - val wrappedHttpClient = - RetryingHttpClient.builder() - .httpClient(perAttemptHttpClient) - .sleeper(sleeper) - .clock(clock) - .maxRetries(maxRetries) - .build() + sharedHttpPipeline + ?: run { + val loggingDelegate = + if ( + requestAuthentication.authenticator != null || + requestAuthentication.attemptAuthenticator != null + ) + httpClient + else + WorkloadIdentityHttpClient( + delegate = httpClient, + workloadIdentityAuth = effectiveWorkloadIdentityAuth, + ) + val loggingHttpClient = + LoggingHttpClient.builder() + .httpClient(loggingDelegate) + .clock(clock) + .level(logLevel) + .propagateAsyncCancellation( + requestAuthentication.attemptAuthenticator != null + ) + .build() + val perAttemptHttpClient = + requestAuthentication.authenticator?.let { authenticator -> + AuthenticatingHttpClient( + delegate = loggingHttpClient, + authenticator = authenticator, + ) + } ?: loggingHttpClient + RetryingHttpClient.builder() + .httpClient(perAttemptHttpClient) + .sleeper(sleeper) + .clock(clock) + .maxRetries(maxRetries) + .attemptAuthenticator(requestAuthentication.attemptAuthenticator) + .build() + } return ClientOptions( httpClient, diff --git a/openai-java-core/src/main/kotlin/com/openai/core/ClientOptionsView.kt b/openai-java-core/src/main/kotlin/com/openai/core/ClientOptionsView.kt new file mode 100644 index 000000000..d87aa73a0 --- /dev/null +++ b/openai-java-core/src/main/kotlin/com/openai/core/ClientOptionsView.kt @@ -0,0 +1,12 @@ +package com.openai.core + +/** + * Adds an SDK presentation header without creating a second owner for an authenticated pipeline. + */ +@JvmSynthetic +internal fun ClientOptions.withDefaultUserAgent(userAgent: String): ClientOptions { + if (headers.values("User-Agent").isNotEmpty()) return this + val builder = toBuilder().putHeader("User-Agent", userAgent) + if (propagatesAsyncCancellation()) builder.shareHttpPipeline(httpClient) + return builder.build() +} diff --git a/openai-java-core/src/main/kotlin/com/openai/core/PrepareRequest.kt b/openai-java-core/src/main/kotlin/com/openai/core/PrepareRequest.kt index 81fee2787..55d4dd962 100644 --- a/openai-java-core/src/main/kotlin/com/openai/core/PrepareRequest.kt +++ b/openai-java-core/src/main/kotlin/com/openai/core/PrepareRequest.kt @@ -5,6 +5,7 @@ package com.openai.core import com.openai.azure.addPathSegmentsForAzure import com.openai.azure.replaceBearerTokenForAzure import com.openai.core.http.HttpRequest +import com.openai.core.http.withPipelineOwnedBody import java.util.Optional import java.util.concurrent.CompletableFuture import kotlin.reflect.full.declaredFunctions @@ -33,10 +34,16 @@ internal fun HttpRequest.prepareAsync( clientOptions: ClientOptions, params: Params, security: SecurityOptions = SecurityOptions.all(), -): CompletableFuture = +): CompletableFuture { // This async version exists to make it easier to add async specific preparation logic in the // future. - CompletableFuture.completedFuture(prepare(clientOptions, params, security)) + val prepared = prepare(clientOptions, params, security) + return if (clientOptions.propagatesAsyncCancellation()) { + CancellationPropagatingFuture.completed(prepared.withPipelineOwnedBody()) + } else { + CompletableFuture.completedFuture(prepared) + } +} @JvmSynthetic internal fun Params.modelNameOrNull(): String? { diff --git a/openai-java-core/src/main/kotlin/com/openai/core/RequestAuthentication.kt b/openai-java-core/src/main/kotlin/com/openai/core/RequestAuthentication.kt index f376b33b2..52cb1812c 100644 --- a/openai-java-core/src/main/kotlin/com/openai/core/RequestAuthentication.kt +++ b/openai-java-core/src/main/kotlin/com/openai/core/RequestAuthentication.kt @@ -1,12 +1,17 @@ package com.openai.core +import com.openai.core.http.HttpRequestAttemptAuthenticator import com.openai.core.http.HttpRequestAuthenticator +import com.openai.core.http.PhantomReachableClosingHttpRequestAttemptAuthenticator import com.openai.core.http.PhantomReachableClosingHttpRequestAuthenticator /** Authentication that is applied to the final HTTP request rather than represented by a key. */ internal sealed interface RequestAuthentication { val authenticator: HttpRequestAuthenticator? + val attemptAuthenticator: HttpRequestAttemptAuthenticator? + get() = null + val fixedBearerBaseUrl: String? get() = null @@ -39,20 +44,21 @@ internal sealed interface RequestAuthentication { class FixedBearerInstalled private constructor( override val fixedBearerBaseUrl: String, - override val authenticator: HttpRequestAuthenticator, + override val attemptAuthenticator: HttpRequestAttemptAuthenticator, ) : RequestAuthentication { - override fun satisfies(security: SecurityOptions): Boolean = - security.bearerAuth && !security.adminApiKeyAuth + override val authenticator: HttpRequestAuthenticator? = null + + override fun satisfies(security: SecurityOptions): Boolean = security.bearerAuth companion object { fun create( fixedBearerBaseUrl: String, - authenticator: HttpRequestAuthenticator, + authenticator: HttpRequestAttemptAuthenticator, ): FixedBearerInstalled = FixedBearerInstalled( fixedBearerBaseUrl, - PhantomReachableClosingHttpRequestAuthenticator(authenticator), + PhantomReachableClosingHttpRequestAttemptAuthenticator(authenticator), ) } } diff --git a/openai-java-core/src/main/kotlin/com/openai/core/RequestOptions.kt b/openai-java-core/src/main/kotlin/com/openai/core/RequestOptions.kt index bf5b28b09..706baf2ab 100644 --- a/openai-java-core/src/main/kotlin/com/openai/core/RequestOptions.kt +++ b/openai-java-core/src/main/kotlin/com/openai/core/RequestOptions.kt @@ -4,6 +4,10 @@ import java.time.Duration class RequestOptions private constructor(val responseValidation: Boolean?, val timeout: Timeout?) { + @JvmSynthetic + internal fun withTimeout(timeout: Timeout): RequestOptions = + RequestOptions(responseValidation, timeout) + companion object { private val NONE = builder().build() diff --git a/openai-java-core/src/main/kotlin/com/openai/core/handlers/ErrorHandler.kt b/openai-java-core/src/main/kotlin/com/openai/core/handlers/ErrorHandler.kt index a235b051a..18f5ec180 100644 --- a/openai-java-core/src/main/kotlin/com/openai/core/handlers/ErrorHandler.kt +++ b/openai-java-core/src/main/kotlin/com/openai/core/handlers/ErrorHandler.kt @@ -10,6 +10,7 @@ import com.openai.core.JsonMissing import com.openai.core.JsonValue import com.openai.core.http.HttpResponse import com.openai.core.http.HttpResponse.Handler +import com.openai.core.http.closeIfPipelineOwned import com.openai.errors.BadRequestException import com.openai.errors.InternalServerException import com.openai.errors.NotFoundException @@ -42,50 +43,57 @@ internal fun errorHandler( errorBodyHandler: Handler> ): Handler = object : Handler { - override fun handle(response: HttpResponse): HttpResponse = - when (val statusCode = response.statusCode()) { - in 200..299 -> response - 400 -> - throw BadRequestException.builder() - .headers(response.headers()) - .error(errorBodyHandler.handle(response)) - .build() - 401 -> - throw UnauthorizedException.builder() - .headers(response.headers()) - .error(errorBodyHandler.handle(response)) - .build() - 403 -> - throw PermissionDeniedException.builder() - .headers(response.headers()) - .error(errorBodyHandler.handle(response)) - .build() - 404 -> - throw NotFoundException.builder() - .headers(response.headers()) - .error(errorBodyHandler.handle(response)) - .build() - 422 -> - throw UnprocessableEntityException.builder() - .headers(response.headers()) - .error(errorBodyHandler.handle(response)) - .build() - 429 -> - throw RateLimitException.builder() - .headers(response.headers()) - .error(errorBodyHandler.handle(response)) - .build() - in 500..599 -> - throw InternalServerException.builder() - .statusCode(statusCode) - .headers(response.headers()) - .error(errorBodyHandler.handle(response)) - .build() - else -> - throw UnexpectedStatusCodeException.builder() - .statusCode(statusCode) - .headers(response.headers()) - .error(errorBodyHandler.handle(response)) - .build() + override fun handle(response: HttpResponse): HttpResponse { + try { + val statusCode = response.statusCode() + if (statusCode in 200..299) return response + when (statusCode) { + 400 -> + throw BadRequestException.builder() + .headers(response.headers()) + .error(errorBodyHandler.handle(response)) + .build() + 401 -> + throw UnauthorizedException.builder() + .headers(response.headers()) + .error(errorBodyHandler.handle(response)) + .build() + 403 -> + throw PermissionDeniedException.builder() + .headers(response.headers()) + .error(errorBodyHandler.handle(response)) + .build() + 404 -> + throw NotFoundException.builder() + .headers(response.headers()) + .error(errorBodyHandler.handle(response)) + .build() + 422 -> + throw UnprocessableEntityException.builder() + .headers(response.headers()) + .error(errorBodyHandler.handle(response)) + .build() + 429 -> + throw RateLimitException.builder() + .headers(response.headers()) + .error(errorBodyHandler.handle(response)) + .build() + in 500..599 -> + throw InternalServerException.builder() + .statusCode(statusCode) + .headers(response.headers()) + .error(errorBodyHandler.handle(response)) + .build() + else -> + throw UnexpectedStatusCodeException.builder() + .statusCode(statusCode) + .headers(response.headers()) + .error(errorBodyHandler.handle(response)) + .build() + } + } catch (failure: Throwable) { + response.closeIfPipelineOwned(failure) + throw failure } + } } diff --git a/openai-java-core/src/main/kotlin/com/openai/core/http/AsyncStreamResponse.kt b/openai-java-core/src/main/kotlin/com/openai/core/http/AsyncStreamResponse.kt index 019dc6aff..7f3d90ab1 100644 --- a/openai-java-core/src/main/kotlin/com/openai/core/http/AsyncStreamResponse.kt +++ b/openai-java-core/src/main/kotlin/com/openai/core/http/AsyncStreamResponse.kt @@ -6,6 +6,9 @@ import java.util.concurrent.CompletableFuture import java.util.concurrent.Executor import java.util.concurrent.atomic.AtomicReference +/** Marker for futures whose cancellation is guaranteed to reach active upstream work. */ +internal interface PropagatesCancellationToUpstream + /** * A class providing access to an API response as an asynchronous stream of chunks of type [T], * where each chunk can be individually processed as soon as it arrives instead of waiting on the @@ -74,7 +77,9 @@ internal fun CompletableFuture>.toAsync(streamHandlerExecu this@toAsync.whenComplete { _, error -> // If an error occurs from the original future, then we should resolve the // `onCompleteFuture` even if `subscribe` has not been called. - error?.let(onCompleteFuture::completeExceptionally) + if (state.get() != State.CLOSED) { + error?.let(onCompleteFuture::completeExceptionally) + } } } @@ -142,6 +147,9 @@ internal fun CompletableFuture>.toAsync(streamHandlerExecu } this@toAsync.whenComplete { streamResponse, error -> streamResponse?.close() } + if (this@toAsync is PropagatesCancellationToUpstream) { + this@toAsync.cancel(true) + } // When the stream is closed, we should always consider it closed. If it closed due // to an error, then we will have already completed the future earlier, and this // will be a no-op. diff --git a/openai-java-core/src/main/kotlin/com/openai/core/http/HttpRequestAttemptAuthenticator.kt b/openai-java-core/src/main/kotlin/com/openai/core/http/HttpRequestAttemptAuthenticator.kt new file mode 100644 index 000000000..eee12f853 --- /dev/null +++ b/openai-java-core/src/main/kotlin/com/openai/core/http/HttpRequestAttemptAuthenticator.kt @@ -0,0 +1,44 @@ +package com.openai.core.http + +import java.time.Duration +import java.util.concurrent.CompletableFuture + +/** + * Reserved authentication seam for integrations that must react to the exact request rejected by + * the server without starting a second retry lifecycle. + */ +interface HttpRequestAttemptAuthenticator : AutoCloseable { + + @JvmSynthetic + fun authenticate(request: HttpRequest, timeout: Duration?): AuthenticatedHttpRequest + + @JvmSynthetic + fun authenticateAsync( + request: HttpRequest, + timeout: Duration?, + ): CompletableFuture = + try { + CompletableFuture.completedFuture(authenticate(request, timeout)) + } catch (throwable: Throwable) { + CompletableFuture().also { + it.completeExceptionally(throwable) + } + } + + override fun close() {} +} + +/** One authenticated wire request and the exact credential generation that produced it. */ +class AuthenticatedHttpRequest +private constructor(private val request: HttpRequest, private val onUnauthorized: () -> Unit) { + + @JvmSynthetic fun request(): HttpRequest = request + + @JvmSynthetic fun onUnauthorized() = onUnauthorized.invoke() + + companion object { + @JvmSynthetic + fun create(request: HttpRequest, onUnauthorized: () -> Unit): AuthenticatedHttpRequest = + AuthenticatedHttpRequest(request, onUnauthorized) + } +} diff --git a/openai-java-core/src/main/kotlin/com/openai/core/http/HttpResponseFor.kt b/openai-java-core/src/main/kotlin/com/openai/core/http/HttpResponseFor.kt index bb1a56199..08a8f1e17 100644 --- a/openai-java-core/src/main/kotlin/com/openai/core/http/HttpResponseFor.kt +++ b/openai-java-core/src/main/kotlin/com/openai/core/http/HttpResponseFor.kt @@ -9,17 +9,34 @@ interface HttpResponseFor : HttpResponse { @JvmSynthetic internal fun HttpResponse.parseable(parse: () -> T): HttpResponseFor = - object : HttpResponseFor { + if (this is PipelineOwnedResource) { + PipelineOwnedHttpResponseFor(this, parse) + } else { + DefaultHttpResponseFor(this, parse) + } - private val parsed: T by lazy { parse() } +private open class DefaultHttpResponseFor(private val response: HttpResponse, parse: () -> T) : + HttpResponseFor { + private val parsed: T by lazy(parse) - override fun parse(): T = parsed + override fun parse(): T = parsed - override fun statusCode(): Int = this@parseable.statusCode() + override fun statusCode(): Int = response.statusCode() - override fun headers(): Headers = this@parseable.headers() + override fun headers(): Headers = response.headers() - override fun body(): InputStream = this@parseable.body() + override fun body(): InputStream = response.body() - override fun close() = this@parseable.close() - } + override fun close() = response.close() +} + +private class PipelineOwnedHttpResponseFor(response: HttpResponse, parse: () -> T) : + DefaultHttpResponseFor(response, parse), PipelineOwnedResource { + override fun parse(): T = + try { + super.parse() + } catch (failure: Throwable) { + closeIfPipelineOwned(failure) + throw failure + } +} diff --git a/openai-java-core/src/main/kotlin/com/openai/core/http/LoggingHttpClient.kt b/openai-java-core/src/main/kotlin/com/openai/core/http/LoggingHttpClient.kt index 57d0d018a..cd51a42dd 100644 --- a/openai-java-core/src/main/kotlin/com/openai/core/http/LoggingHttpClient.kt +++ b/openai-java-core/src/main/kotlin/com/openai/core/http/LoggingHttpClient.kt @@ -19,6 +19,7 @@ import java.time.OffsetDateTime import java.util.SortedSet import java.util.concurrent.CompletableFuture import java.util.concurrent.CompletionException +import java.util.concurrent.atomic.AtomicBoolean import kotlin.time.toKotlinDuration /** A wrapper [HttpClient] around [httpClient] that logs request and response information. */ @@ -46,6 +47,7 @@ private constructor( * Pass [LogLevel.fromEnv] to read from environment variables. */ @get:JvmName("level") val level: LogLevel, + private val propagateAsyncCancellation: Boolean, ) : HttpClient { override fun execute(request: HttpRequest, requestOptions: RequestOptions): HttpResponse { @@ -60,8 +62,18 @@ private constructor( throw e } - val took = Duration.between(before, OffsetDateTime.now(clock)) - return logResponse(response, took) + if (!propagateAsyncCancellation) { + val took = Duration.between(before, OffsetDateTime.now(clock)) + return logResponse(response, took) + } + val lease = PipelineResponseLease() + return try { + val took = Duration.between(before, OffsetDateTime.now(clock)) + logResponse(lease.acquire(response), took) + } catch (failure: Throwable) { + lease.close(response, failure) + throw failure + } } override fun executeAsync( @@ -78,6 +90,9 @@ private constructor( logFailure(e, Duration.between(before, OffsetDateTime.now(clock))) throw e } + if (propagateAsyncCancellation) { + return loggingFutureWithCancellationPropagation(future, before) + } return future.handle { response, error -> val took = Duration.between(before, OffsetDateTime.now(clock)) if (error != null) { @@ -88,6 +103,44 @@ private constructor( } } + private fun loggingFutureWithCancellationPropagation( + source: CompletableFuture, + before: OffsetDateTime, + ): CompletableFuture { + val lease = PipelineResponseLease() + val result = CompletableFuture() + + fun discard(response: HttpResponse, failure: Throwable? = null) { + if (failure == null) lease.discard(response) else lease.close(response, failure) + } + + source.whenComplete { response, error -> + try { + if (result.isCancelled) { + if (error == null) discard(response) + return@whenComplete + } + val took = Duration.between(before, OffsetDateTime.now(clock)) + if (error != null) { + logFailure(unwrapCompletionException(error), took) + result.completeExceptionally(error) + return@whenComplete + } + val logged = logResponse(lease.acquire(response), took) + if (!result.complete(logged)) discard(logged) + } catch (failure: Throwable) { + if (response != null) discard(response, failure) + result.completeExceptionally(failure) + } + } + result.whenComplete { _, _ -> + if (result.isCancelled && !source.cancel(true)) { + source.whenComplete { response, error -> if (error == null) discard(response) } + } + } + return result + } + private fun logRequest(request: HttpRequest): HttpRequest { if (!level.shouldLog(LogLevel.INFO)) { return request @@ -139,7 +192,11 @@ private constructor( } logHeaders(response.headers()) - return LoggingHttpResponse(response) + return if (response is PipelineOwnedResource) { + PipelineOwnedLoggingHttpResponse(response) + } else { + LoggingHttpResponse(response) + } } private fun logFailure(error: Throwable, took: Duration) { @@ -199,6 +256,7 @@ private constructor( ) private var clock: Clock = Clock.systemUTC() private var level: LogLevel? = null + private var propagateAsyncCancellation: Boolean = false @JvmSynthetic internal fun from(loggingHttpClient: LoggingHttpClient) = apply { @@ -206,6 +264,7 @@ private constructor( redactedHeaders = loggingHttpClient.redactedHeaders clock = loggingHttpClient.clock level = loggingHttpClient.level + propagateAsyncCancellation = loggingHttpClient.propagateAsyncCancellation } /** The underlying [HttpClient] for making requests. */ @@ -237,6 +296,11 @@ private constructor( */ fun level(level: LogLevel) = apply { this.level = level } + @JvmSynthetic + internal fun propagateAsyncCancellation(propagate: Boolean) = apply { + propagateAsyncCancellation = propagate + } + /** * Returns an immutable instance of [LoggingHttpClient]. * @@ -256,6 +320,7 @@ private constructor( redactedHeaders.toSortedSet(String.CASE_INSENSITIVE_ORDER).toImmutable(), clock, checkRequired("level", level), + propagateAsyncCancellation, ) } } @@ -332,7 +397,7 @@ private class LoggingOutputStream(private val outputStream: OutputStream, charse * * The logging occurs in a streaming manner with minimal buffering. */ -private class LoggingHttpResponse(private val response: HttpResponse) : HttpResponse { +private open class LoggingHttpResponse(private val response: HttpResponse) : HttpResponse { private val loggingBody: Lazy = lazy { LoggingInputStream( @@ -347,11 +412,29 @@ private class LoggingHttpResponse(private val response: HttpResponse) : HttpResp override fun body(): InputStream = loggingBody.value - override fun close() { - if (loggingBody.isInitialized()) { - loggingBody.value.close() + open override fun close() { + var failure: Throwable? = null + try { + if (loggingBody.isInitialized()) loggingBody.value.close() + } catch (error: Throwable) { + failure = error + } + try { + response.close() + } catch (error: Throwable) { + if (failure == null) failure = error + else if (failure !== error) failure.addSuppressed(error) } - response.close() + failure?.let { throw it } + } +} + +private class PipelineOwnedLoggingHttpResponse(response: HttpResponse) : + LoggingHttpResponse(response), PipelineOwnedResource { + private val closed = AtomicBoolean() + + override fun close() { + if (closed.compareAndSet(false, true)) super.close() } } diff --git a/openai-java-core/src/main/kotlin/com/openai/core/http/PhantomReachableClosingHttpRequestAttemptAuthenticator.kt b/openai-java-core/src/main/kotlin/com/openai/core/http/PhantomReachableClosingHttpRequestAttemptAuthenticator.kt new file mode 100644 index 000000000..fc62de3e3 --- /dev/null +++ b/openai-java-core/src/main/kotlin/com/openai/core/http/PhantomReachableClosingHttpRequestAttemptAuthenticator.kt @@ -0,0 +1,25 @@ +package com.openai.core.http + +import com.openai.core.closeWhenPhantomReachable +import java.time.Duration +import java.util.concurrent.CompletableFuture + +/** Closes a delegated attempt authenticator after this wrapper becomes phantom reachable. */ +internal class PhantomReachableClosingHttpRequestAttemptAuthenticator( + private val authenticator: HttpRequestAttemptAuthenticator +) : HttpRequestAttemptAuthenticator { + init { + closeWhenPhantomReachable(this, authenticator) + } + + override fun authenticate(request: HttpRequest, timeout: Duration?): AuthenticatedHttpRequest = + authenticator.authenticate(request, timeout) + + override fun authenticateAsync( + request: HttpRequest, + timeout: Duration?, + ): CompletableFuture = + authenticator.authenticateAsync(request, timeout) + + override fun close() = authenticator.close() +} diff --git a/openai-java-core/src/main/kotlin/com/openai/core/http/PipelineRequestBody.kt b/openai-java-core/src/main/kotlin/com/openai/core/http/PipelineRequestBody.kt new file mode 100644 index 000000000..e5002c66b --- /dev/null +++ b/openai-java-core/src/main/kotlin/com/openai/core/http/PipelineRequestBody.kt @@ -0,0 +1,52 @@ +package com.openai.core.http + +import java.io.OutputStream +import java.util.concurrent.atomic.AtomicBoolean + +/** Request body whose terminal cleanup is owned by the authenticated request pipeline. */ +internal interface PipelineOwnedRequestBody : HttpRequestBody + +/** Shares one close across pre-dispatch cancellation, the transport, and terminal cleanup. */ +@JvmSynthetic +internal fun HttpRequest.withPipelineOwnedBody(): HttpRequest { + val current = body ?: return this + if (current is PipelineOwnedRequestBody) return this + return toBuilder().body(CloseOncePipelineRequestBody(current)).build() +} + +/** Best-effort terminal cleanup for an authenticated request body. */ +@JvmSynthetic +internal fun HttpRequest.closePipelineBody(failure: Throwable? = null) { + val owned = body as? PipelineOwnedRequestBody ?: return + try { + owned.close() + } catch (closeFailure: Throwable) { + if (failure == null) throw closeFailure + if (closeFailure !== failure) failure.addSuppressed(closeFailure) + } +} + +/** Best-effort terminal cleanup for cancellation or an already completed async result. */ +@JvmSynthetic +internal fun HttpRequest.discardPipelineBody() { + try { + closePipelineBody() + } catch (_: Throwable) {} +} + +private class CloseOncePipelineRequestBody(private val delegate: HttpRequestBody) : + PipelineOwnedRequestBody { + private val closed = AtomicBoolean() + + override fun writeTo(outputStream: OutputStream) = delegate.writeTo(outputStream) + + override fun contentType(): String? = delegate.contentType() + + override fun contentLength(): Long = delegate.contentLength() + + override fun repeatable(): Boolean = delegate.repeatable() + + override fun close() { + if (closed.compareAndSet(false, true)) delegate.close() + } +} diff --git a/openai-java-core/src/main/kotlin/com/openai/core/http/PipelineResponseLease.kt b/openai-java-core/src/main/kotlin/com/openai/core/http/PipelineResponseLease.kt new file mode 100644 index 000000000..505a61978 --- /dev/null +++ b/openai-java-core/src/main/kotlin/com/openai/core/http/PipelineResponseLease.kt @@ -0,0 +1,63 @@ +package com.openai.core.http + +import java.util.concurrent.atomic.AtomicBoolean +import java.util.concurrent.atomic.AtomicReference + +/** Resource whose lifecycle is owned by the authenticated request pipeline. */ +internal interface PipelineOwnedResource : AutoCloseable + +/** Closes a pipeline-owned resource and preserves a primary failure when one exists. */ +@JvmSynthetic +internal fun Any?.closeIfPipelineOwned(failure: Throwable? = null) { + val resource = this as? PipelineOwnedResource ?: return + try { + resource.close() + } catch (closeFailure: Throwable) { + if (failure == null) throw closeFailure + if (closeFailure !== failure) failure.addSuppressed(closeFailure) + } +} + +/** Returns one close-once pipeline wrapper for a response. */ +@JvmSynthetic +internal fun HttpResponse.asPipelineOwned(): HttpResponse = + if (this is PipelineOwnedResource) this else CloseOncePipelineHttpResponse(this) + +/** Shares one response wrapper between completion and cancellation race participants. */ +internal class PipelineResponseLease { + private val response = AtomicReference() + + fun acquire(candidate: HttpResponse): HttpResponse { + response.get()?.let { + return it + } + val owned = candidate.asPipelineOwned() + return if (response.compareAndSet(null, owned)) owned else requireNotNull(response.get()) + } + + fun close(candidate: HttpResponse, failure: Throwable? = null) { + acquire(candidate).closeIfPipelineOwned(failure) + } + + fun discard(candidate: HttpResponse) { + try { + close(candidate) + } catch (_: Throwable) {} + } +} + +/** Makes cancellation-time and parser-finally cleanup one underlying response close. */ +internal class CloseOncePipelineHttpResponse(private val delegate: HttpResponse) : + HttpResponse, PipelineOwnedResource { + private val closed = AtomicBoolean() + + override fun statusCode(): Int = delegate.statusCode() + + override fun headers(): Headers = delegate.headers() + + override fun body() = delegate.body() + + override fun close() { + if (closed.compareAndSet(false, true)) delegate.close() + } +} diff --git a/openai-java-core/src/main/kotlin/com/openai/core/http/RetryingHttpClient.kt b/openai-java-core/src/main/kotlin/com/openai/core/http/RetryingHttpClient.kt index 317205674..38c9167a6 100644 --- a/openai-java-core/src/main/kotlin/com/openai/core/http/RetryingHttpClient.kt +++ b/openai-java-core/src/main/kotlin/com/openai/core/http/RetryingHttpClient.kt @@ -1,229 +1,27 @@ package com.openai.core.http import com.openai.core.DefaultSleeper -import com.openai.core.RequestOptions import com.openai.core.Sleeper import com.openai.core.checkRequired -import com.openai.errors.OpenAIIoException -import com.openai.errors.OpenAIRetryableException -import java.io.IOException import java.time.Clock -import java.time.Duration -import java.time.OffsetDateTime -import java.time.format.DateTimeFormatter -import java.time.format.DateTimeParseException -import java.time.temporal.ChronoUnit -import java.util.UUID -import java.util.concurrent.CompletableFuture -import java.util.concurrent.ThreadLocalRandom -import java.util.concurrent.TimeUnit -import java.util.function.Function -import kotlin.math.min -import kotlin.math.pow class RetryingHttpClient private constructor( - private val httpClient: HttpClient, - private val sleeper: Sleeper, - private val clock: Clock, - private val maxRetries: Int, - private val idempotencyHeader: String?, -) : HttpClient { - - override fun execute(request: HttpRequest, requestOptions: RequestOptions): HttpResponse { - var modifiedRequest = maybeAddIdempotencyHeader(request) - - // Don't send the current retry count in the headers if the caller set their own value. - val shouldSendRetryCount = - !modifiedRequest.headers.names().contains("X-Stainless-Retry-Count") - - var retries = 0 - - while (true) { - if (shouldSendRetryCount) { - modifiedRequest = setRetryCountHeader(modifiedRequest, retries) - } - - if (!isRetryable(modifiedRequest)) { - return httpClient.execute(modifiedRequest, requestOptions) - } - - val response = - try { - val response = httpClient.execute(modifiedRequest, requestOptions) - if (++retries > maxRetries || !shouldRetry(response)) { - return response - } - - response - } catch (throwable: Throwable) { - if (++retries > maxRetries || !shouldRetry(throwable)) { - throw throwable - } - - null - } - - val backoffDuration = getRetryBackoffDuration(retries, response) - // All responses must be closed, so close the failed one before retrying. - response?.close() - sleeper.sleep(backoffDuration) - } - } - - override fun executeAsync( - request: HttpRequest, - requestOptions: RequestOptions, - ): CompletableFuture { - val modifiedRequest = maybeAddIdempotencyHeader(request) - - // Don't send the current retry count in the headers if the caller set their own value. - val shouldSendRetryCount = - !modifiedRequest.headers.names().contains("X-Stainless-Retry-Count") - - var retries = 0 - - fun executeWithRetries( - request: HttpRequest, - requestOptions: RequestOptions, - ): CompletableFuture { - val requestWithRetryCount = - if (shouldSendRetryCount) setRetryCountHeader(request, retries) else request - - val responseFuture = httpClient.executeAsync(requestWithRetryCount, requestOptions) - if (!isRetryable(requestWithRetryCount)) { - return responseFuture - } - - return responseFuture - .handleAsync( - fun( - response: HttpResponse?, - throwable: Throwable?, - ): CompletableFuture { - if (response != null) { - if (++retries > maxRetries || !shouldRetry(response)) { - return CompletableFuture.completedFuture(response) - } - } else { - if (++retries > maxRetries || !shouldRetry(throwable!!)) { - val failedFuture = CompletableFuture() - failedFuture.completeExceptionally(throwable) - return failedFuture - } - } - - val backoffDuration = getRetryBackoffDuration(retries, response) - // All responses must be closed, so close the failed one before retrying. - response?.close() - return sleeper.sleepAsync(backoffDuration).thenCompose { - executeWithRetries(requestWithRetryCount, requestOptions) - } - } - ) { - // Run in the same thread. - it.run() - } - .thenCompose(Function.identity()) - } - - return executeWithRetries(modifiedRequest, requestOptions) - } - - override fun close() { - httpClient.close() - sleeper.close() - } - - private fun isRetryable(request: HttpRequest): Boolean = - // Some requests, such as when a request body is being streamed, cannot be retried because - // the body data aren't available on subsequent attempts. - request.body?.repeatable() ?: true - - private fun setRetryCountHeader(request: HttpRequest, retries: Int): HttpRequest = - request.toBuilder().replaceHeaders("X-Stainless-Retry-Count", retries.toString()).build() - - private fun idempotencyKey(): String = "stainless-java-retry-${UUID.randomUUID()}" - - private fun maybeAddIdempotencyHeader(request: HttpRequest): HttpRequest { - if (idempotencyHeader == null || request.headers.names().contains(idempotencyHeader)) { - return request - } - - return request - .toBuilder() - // Set a header to uniquely identify the request when retried. - .putHeader(idempotencyHeader, idempotencyKey()) - .build() - } - - private fun shouldRetry(response: HttpResponse): Boolean { - // Note: this is not a standard header - val shouldRetryHeader = response.headers().values("X-Should-Retry").getOrNull(0) - val statusCode = response.statusCode() - - return when { - // If the server explicitly says whether to retry, obey - shouldRetryHeader == "true" -> true - shouldRetryHeader == "false" -> false - - // Retry on request timeouts - statusCode == 408 -> true - // Retry on lock timeouts - statusCode == 409 -> true - // Retry on rate limits - statusCode == 429 -> true - // Retry internal errors - statusCode >= 500 -> true - else -> false - } - } - - private fun shouldRetry(throwable: Throwable): Boolean = - // Only retry known retryable exceptions, other exceptions are not intended to be retried. - throwable is IOException || - throwable is OpenAIIoException || - throwable is OpenAIRetryableException - - private fun getRetryBackoffDuration(retries: Int, response: HttpResponse?): Duration { - // About the Retry-After header: - // https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Retry-After - response - ?.headers() - ?.let { headers -> - headers - .values("Retry-After-Ms") - .getOrNull(0) - ?.toFloatOrNull() - ?.times(TimeUnit.MILLISECONDS.toNanos(1)) - ?: headers.values("Retry-After").getOrNull(0)?.let { retryAfter -> - retryAfter.toFloatOrNull()?.times(TimeUnit.SECONDS.toNanos(1)) - ?: try { - ChronoUnit.NANOS.between( - OffsetDateTime.now(clock), - OffsetDateTime.parse( - retryAfter, - DateTimeFormatter.RFC_1123_DATE_TIME, - ), - ) - } catch (e: DateTimeParseException) { - null - } - } - } - ?.let { retryAfterNanos -> - // If the API asks us to wait a certain amount of time, do what it says. - return Duration.ofNanos(retryAfterNanos.toLong()) - } - - // Apply exponential backoff, but not more than the max. - val backoffSeconds = min(0.5 * 2.0.pow(retries - 1), 8.0) - - // Apply some jitter - val jitter = 1.0 - 0.25 * ThreadLocalRandom.current().nextDouble() - - return Duration.ofNanos((TimeUnit.SECONDS.toNanos(1) * backoffSeconds * jitter).toLong()) - } + httpClient: HttpClient, + sleeper: Sleeper, + clock: Clock, + maxRetries: Int, + idempotencyHeader: String?, + attemptAuthenticator: HttpRequestAttemptAuthenticator?, +) : + HttpClient by RetryingHttpClientOrchestrator( + httpClient, + sleeper, + clock, + maxRetries, + idempotencyHeader, + attemptAuthenticator, + ) { companion object { @@ -237,6 +35,7 @@ private constructor( private var clock: Clock = Clock.systemUTC() private var maxRetries: Int = 2 private var idempotencyHeader: String? = null + private var attemptAuthenticator: HttpRequestAttemptAuthenticator? = null fun httpClient(httpClient: HttpClient) = apply { this.httpClient = httpClient } @@ -248,6 +47,11 @@ private constructor( fun idempotencyHeader(header: String) = apply { this.idempotencyHeader = header } + @JvmSynthetic + fun attemptAuthenticator(authenticator: HttpRequestAttemptAuthenticator?) = apply { + this.attemptAuthenticator = authenticator + } + fun build(): HttpClient = RetryingHttpClient( checkRequired("httpClient", httpClient), @@ -255,6 +59,7 @@ private constructor( clock, maxRetries, idempotencyHeader, + attemptAuthenticator, ) } } diff --git a/openai-java-core/src/main/kotlin/com/openai/core/http/RetryingHttpClientOrchestrator.kt b/openai-java-core/src/main/kotlin/com/openai/core/http/RetryingHttpClientOrchestrator.kt new file mode 100644 index 000000000..eaf6f6e2c --- /dev/null +++ b/openai-java-core/src/main/kotlin/com/openai/core/http/RetryingHttpClientOrchestrator.kt @@ -0,0 +1,666 @@ +package com.openai.core.http + +import com.openai.core.RequestOptions +import com.openai.core.Sleeper +import com.openai.core.Timeout +import com.openai.errors.OpenAIIoException +import com.openai.errors.OpenAIRetryableException +import com.openai.errors.UnexpectedStatusCodeException +import java.io.IOException +import java.time.Clock +import java.time.Duration +import java.time.OffsetDateTime +import java.time.format.DateTimeFormatter +import java.time.format.DateTimeParseException +import java.time.temporal.ChronoUnit +import java.util.UUID +import java.util.concurrent.CompletableFuture +import java.util.concurrent.CompletionException +import java.util.concurrent.ConcurrentHashMap +import java.util.concurrent.ExecutionException +import java.util.concurrent.ThreadLocalRandom +import java.util.concurrent.TimeUnit +import java.util.concurrent.atomic.AtomicBoolean +import java.util.concurrent.atomic.AtomicReference +import java.util.function.Function +import kotlin.math.min +import kotlin.math.pow + +internal class RetryingHttpClientOrchestrator( + private val httpClient: HttpClient, + private val sleeper: Sleeper, + private val clock: Clock, + private val maxRetries: Int, + private val idempotencyHeader: String?, + private val attemptAuthenticator: HttpRequestAttemptAuthenticator?, + private val nanoTime: () -> Long = System::nanoTime, +) : HttpClient { + private val closed = AtomicBoolean() + private val activeAuthenticatedRequests = ConcurrentHashMap.newKeySet>() + private var authenticatorClosed = attemptAuthenticator == null + private var httpClientClosed = false + private var sleeperClosed = false + + private fun trackAuthenticatedRequest(future: CompletableFuture<*>): Boolean { + if (closed.get()) return false + activeAuthenticatedRequests.add(future) + future.whenComplete { _, _ -> activeAuthenticatedRequests.remove(future) } + if (closed.get() && activeAuthenticatedRequests.remove(future)) { + future.completeExceptionally(OpenAIIoException("HTTP client is closed")) + return false + } + return true + } + + override fun execute(request: HttpRequest, requestOptions: RequestOptions): HttpResponse = + executeWithRetries(request, requestOptions) + + override fun executeAsync( + request: HttpRequest, + requestOptions: RequestOptions, + ): CompletableFuture = executeWithRetriesAsync(request, requestOptions) + + @Synchronized + override fun close() { + closed.set(true) + activeAuthenticatedRequests.forEach { future -> + future.completeExceptionally(OpenAIIoException("HTTP client is closed")) + } + if (attemptAuthenticator == null) { + httpClient.close() + sleeper.close() + return + } + var failure: Throwable? = null + fun closeComponent(isClosed: Boolean, close: () -> Unit): Boolean { + if (isClosed) return true + try { + close() + return true + } catch (error: Throwable) { + if (failure == null) failure = error + else if (error !== failure) failure?.addSuppressed(error) + return false + } + } + authenticatorClosed = closeComponent(authenticatorClosed, attemptAuthenticator::close) + httpClientClosed = closeComponent(httpClientClosed, httpClient::close) + sleeperClosed = closeComponent(sleeperClosed, sleeper::close) + failure?.let { throw it } + } + + private fun executeWithRetries( + request: HttpRequest, + requestOptions: RequestOptions, + ): HttpResponse { + val authenticator = attemptAuthenticator + val pipelineRequest = + if (authenticator != null) request.withPipelineOwnedBody() else request + var requestFailure: Throwable? = null + try { + if (authenticator != null && closed.get()) { + throw OpenAIIoException("HTTP client is closed") + } + val deadline = authenticator?.let { deadline(requestOptions) } + val modified = maybeAddIdempotencyHeader(pipelineRequest) + val sendRetryCount = !modified.headers.names().contains("X-Stainless-Retry-Count") + var retries = 0 + var replayed = false + while (true) { + val current = + if (sendRetryCount) setRetryCountHeader(modified, retries) else modified + if (authenticator == null) { + if (!isRetryable(current)) return httpClient.execute(current, requestOptions) + val response = + try { + val value = httpClient.execute(current, requestOptions) + if (++retries > maxRetries || !shouldRetry(value)) return value + value + } catch (error: Throwable) { + if (++retries > maxRetries || !shouldRetry(error)) throw error + null + } + val delay = getRetryBackoffDuration(retries, response) + response?.close() + sleeper.sleep(delay) + continue + } + val authenticated = + try { + authenticator.authenticate(current, remainingOrThrow(deadline)) + } catch (error: Throwable) { + if (retries >= maxRetries || !shouldRetryAttempt(error)) { + throw error + } + retries++ + sleepAuthenticated( + getRetryBackoffDuration(retries, error = error), + deadline, + ) + continue + } + val options = + deadline?.let { remainingOptions(requestOptions, it) } ?: requestOptions + val authenticatedRequest = authenticated.request() + val response = + try { + httpClient.execute(authenticatedRequest, options) + } catch (error: Throwable) { + if ( + !isRetryable(authenticatedRequest) || + retries >= maxRetries || + !shouldRetryAttempt(error) + ) { + throw error + } + retries++ + sleepAuthenticated(getRetryBackoffDuration(retries), deadline) + continue + } + var ownedResponse: HttpResponse? = response + try { + if (response.statusCode() == 401) { + authenticated.onUnauthorized() + // Authentication replay is the separate `replayUsed` bit in RequestScope. + // It + // shares the deadline and retry-count continuity, but does not consume or + // reset + // the transient maxRetries budget. + if (isRetryable(authenticatedRequest) && !replayed) { + replayed = true + ownedResponse = null + response.close() + continue + } + ownedResponse = null + return response.asPipelineOwned() + } + if ( + !isRetryable(authenticatedRequest) || + retries >= maxRetries || + !shouldRetry(response) + ) { + ownedResponse = null + return response.asPipelineOwned() + } + retries++ + val delay = getRetryBackoffDuration(retries, response) + ownedResponse = null + response.close() + sleepAuthenticated(delay, deadline) + } catch (failure: Throwable) { + ownedResponse?.let { + try { + it.close() + } catch (closeFailure: Throwable) { + if (closeFailure !== failure) failure.addSuppressed(closeFailure) + } + } + throw failure + } + } + } catch (failure: Throwable) { + requestFailure = failure + throw failure + } finally { + if (authenticator != null) pipelineRequest.closePipelineBody(requestFailure) + } + } + + private fun executeWithRetriesAsync( + request: HttpRequest, + requestOptions: RequestOptions, + ): CompletableFuture = + attemptAuthenticator?.let { + executeAuthenticatedWithRetriesAsync(request, requestOptions, it) + } ?: executeOrdinaryWithRetriesAsync(request, requestOptions) + + private fun executeAuthenticatedWithRetriesAsync( + request: HttpRequest, + requestOptions: RequestOptions, + authenticator: HttpRequestAttemptAuthenticator, + ): CompletableFuture { + val pipelineRequest = request.withPipelineOwnedBody() + val result = CompletableFuture() + val active = AtomicReference?>() + result.whenComplete { _, _ -> + if (result.isCancelled || closed.get()) active.getAndSet(null)?.cancel(true) + pipelineRequest.discardPipelineBody() + } + if (closed.get()) { + result.completeExceptionally(OpenAIIoException("HTTP client is closed")) + return result + } + if (!trackAuthenticatedRequest(result)) { + result.completeExceptionally(OpenAIIoException("HTTP client is closed")) + return result + } + val deadline = deadline(requestOptions) + val modified = maybeAddIdempotencyHeader(pipelineRequest) + val sendRetryCount = !modified.headers.names().contains("X-Stainless-Retry-Count") + var retries = 0 + var replayed = false + + fun activate(future: CompletableFuture<*>) { + active.set(future) + if (result.isDone) future.cancel(true) + } + + fun closeDiscarded(response: HttpResponse?) { + try { + response?.close() + } catch (_: Throwable) {} + } + + fun fail(error: Throwable, response: HttpResponse? = null) { + var failure = error + if (response != null) { + try { + response.close() + } catch (closeError: Throwable) { + if (closeError !== failure) failure.addSuppressed(closeError) + } + } + result.completeExceptionally(failure) + } + + fun completeResponse(response: HttpResponse) { + val delivered = response.asPipelineOwned() + if (!result.complete(delivered)) delivered.close() + } + + lateinit var run: () -> Unit + fun retry( + response: HttpResponse? = null, + error: Throwable? = null, + authenticationFailure: Boolean = false, + requestRetryable: Boolean = true, + ) { + var ownedResponse = response + try { + if (result.isDone) { + ownedResponse = null + closeDiscarded(response) + return + } + if (!requestRetryable && !authenticationFailure) { + if (response != null) { + ownedResponse = null + completeResponse(response) + } else { + result.completeExceptionally( + error ?: IllegalStateException("Missing failure") + ) + } + return + } + val cause = error?.let(::unwrap) + val retryError = cause?.let(::shouldRetryAttempt) + if (retries >= maxRetries || (response?.let(::shouldRetry) ?: retryError != true)) { + if (response != null) { + ownedResponse = null + completeResponse(response) + } else { + result.completeExceptionally( + cause ?: IllegalStateException("Missing failure") + ) + } + return + } + retries++ + val delay = getRetryBackoffDuration(retries, response, cause) + if (response != null) { + ownedResponse = null + response.close() + } + val remaining = deadline?.let(::remaining) + if (remaining != null && delay >= remaining) { + result.completeExceptionally(timedOut()) + return + } + val sleep = + try { + sleeper.sleepAsync(delay) + } catch (sleepError: Throwable) { + result.completeExceptionally(sleepError) + return + } + activate(sleep) + sleep.whenComplete { _, sleepError -> + try { + if (sleepError == null) run() + else { + result.completeExceptionally(unwrap(sleepError)) + } + } catch (sleepFailure: Throwable) { + fail(sleepFailure) + } + } + } catch (retryFailure: Throwable) { + fail(retryFailure, ownedResponse) + } + } + + fun dispatch(authenticated: AuthenticatedHttpRequest, options: RequestOptions) { + val authenticatedRequest = authenticated.request() + val call = + try { + httpClient.executeAsync(authenticatedRequest, options) + } catch (error: Throwable) { + retry(error = error, requestRetryable = isRetryable(authenticatedRequest)) + return + } + activate(call) + call.whenComplete callComplete@{ response, callError -> + var ownedResponse = response + try { + if (result.isDone) { + ownedResponse = null + closeDiscarded(response) + return@callComplete + } + val requestRetryable = isRetryable(authenticatedRequest) + if (callError != null) { + ownedResponse = null + closeDiscarded(response) + retry(error = callError, requestRetryable = requestRetryable) + } else if (response.statusCode() == 401) { + authenticated.onUnauthorized() + // See the synchronous path: replayUsed is distinct from maxRetries. + if (requestRetryable && !replayed) { + replayed = true + ownedResponse = null + response.close() + run() + } else { + ownedResponse = null + completeResponse(response) + } + } else { + ownedResponse = null + retry(response = response, requestRetryable = requestRetryable) + } + } catch (callbackFailure: Throwable) { + fail(callbackFailure, ownedResponse) + } + } + } + + run = run@{ + if (result.isDone) return@run + val timeout = + try { + remainingOrThrow(deadline) + } catch (error: Throwable) { + result.completeExceptionally(error) + return@run + } + val current = if (sendRetryCount) setRetryCountHeader(modified, retries) else modified + val authentication = + try { + authenticator.authenticateAsync(current, timeout) + } catch (error: Throwable) { + retry(error = error, authenticationFailure = true) + return@run + } + activate(authentication) + authentication.whenComplete authenticationComplete@{ authenticated, authError -> + try { + if (result.isDone) return@authenticationComplete + if (authError != null) { + retry(error = authError, authenticationFailure = true) + return@authenticationComplete + } + val options = + try { + deadline?.let { remainingOptions(requestOptions, it) } ?: requestOptions + } catch (error: Throwable) { + result.completeExceptionally(error) + return@authenticationComplete + } + dispatch(authenticated, options) + } catch (error: Throwable) { + fail(error) + } + } + } + run() + return result + } + + /** + * The pre-X.509 pipeline is kept intact so ordinary-client completion semantics do not move. + */ + private fun executeOrdinaryWithRetriesAsync( + request: HttpRequest, + requestOptions: RequestOptions, + ): CompletableFuture { + val modifiedRequest = maybeAddIdempotencyHeader(request) + val shouldSendRetryCount = + !modifiedRequest.headers.names().contains("X-Stainless-Retry-Count") + var retries = 0 + + fun executeWithRetries(current: HttpRequest): CompletableFuture { + val requestWithRetryCount = + if (shouldSendRetryCount) setRetryCountHeader(current, retries) else current + val responseFuture = httpClient.executeAsync(requestWithRetryCount, requestOptions) + if (!isRetryable(requestWithRetryCount)) return responseFuture + + return responseFuture + .handleAsync( + fun( + response: HttpResponse?, + error: Throwable?, + ): CompletableFuture { + if (response != null) { + if (++retries > maxRetries || !shouldRetry(response)) { + return CompletableFuture.completedFuture(response) + } + } else if (++retries > maxRetries || !shouldRetry(error!!)) { + return CompletableFuture().also { + it.completeExceptionally(error) + } + } + + val backoffDuration = getRetryBackoffDuration(retries, response) + response?.close() + return sleeper.sleepAsync(backoffDuration).thenCompose { + executeWithRetries(requestWithRetryCount) + } + } + ) { + it.run() + } + .thenCompose(Function.identity()) + } + + return executeWithRetries(modifiedRequest) + } + + private fun deadline(options: RequestOptions): Deadline? { + val duration = (options.timeout ?: Timeout.default()).request() + if (duration.isZero) return null + val nanos = + try { + duration.toNanos() + } catch (_: ArithmeticException) { + if (duration.isNegative) 0 else Long.MAX_VALUE + } + return Deadline(nanoTime(), nanos) + } + + private fun remaining(deadline: Deadline): Duration { + val elapsed = nanoTime() - deadline.startedAt + val nanos = + if (elapsed < 0 || elapsed >= deadline.timeoutNanos) 0 + else deadline.timeoutNanos - elapsed + return Duration.ofNanos(nanos) + } + + private fun remainingOrThrow(deadline: Deadline?): Duration? = + deadline?.let(::remaining)?.also { if (it.isZero) throw timedOut() } + + private fun remainingOptions(options: RequestOptions, deadline: Deadline): RequestOptions { + val remaining = remaining(deadline) + if (remaining.isZero) throw timedOut() + return options.withTimeout( + (options.timeout ?: Timeout.default()).toBuilder().request(remaining).build() + ) + } + + private fun sleepAuthenticated(delay: Duration, deadline: Deadline?) { + val remaining = deadline?.let(::remaining) + if (remaining != null && delay >= remaining) throw timedOut() + val closeWaiter = CompletableFuture() + if (!trackAuthenticatedRequest(closeWaiter)) { + throw OpenAIIoException("HTTP client is closed") + } + if (closeWaiter.isDone) { + throw OpenAIIoException("HTTP client is closed") + } + if (closed.get()) { + closeWaiter.cancel(false) + throw OpenAIIoException("HTTP client is closed") + } + val sleep = + try { + sleeper.sleepAsync(delay) + } catch (error: Throwable) { + closeWaiter.cancel(false) + throw error + } + if (sleep.isDone && !sleep.isCompletedExceptionally) { + closeWaiter.cancel(false) + sleep.get() + return + } + try { + CompletableFuture.anyOf(sleep, closeWaiter).get() + if (closeWaiter.isDone) throw OpenAIIoException("HTTP client is closed") + } catch (error: InterruptedException) { + Thread.currentThread().interrupt() + throw OpenAIIoException("Interrupted during retry backoff", error) + } catch (error: ExecutionException) { + val cause = unwrap(error) + if (cause is OpenAIIoException && cause.message == "HTTP client is closed") { + throw cause + } + throw OpenAIIoException("Retry backoff failed", cause) + } finally { + closeWaiter.cancel(false) + sleep.cancel(true) + } + } + + private fun timedOut() = OpenAIIoException("X.509 request deadline exceeded") + + private data class Deadline(val startedAt: Long, val timeoutNanos: Long) + + private fun shouldRetryAttempt(error: Throwable): Boolean { + val cause = unwrap(error) + return if (cause is UnexpectedStatusCodeException) { + shouldRetry(cause.statusCode(), cause.headers()) + } else shouldRetry(cause) + } + + private fun unwrap(error: Throwable): Throwable = + if (error is CompletionException || error is ExecutionException) error.cause ?: error + else error + + private fun isRetryable(request: HttpRequest): Boolean = + // Some requests, such as when a request body is being streamed, cannot be retried because + // the body data aren't available on subsequent attempts. + request.body?.repeatable() ?: true + + private fun setRetryCountHeader(request: HttpRequest, retries: Int): HttpRequest = + request.toBuilder().replaceHeaders("X-Stainless-Retry-Count", retries.toString()).build() + + private fun idempotencyKey(): String = "stainless-java-retry-${UUID.randomUUID()}" + + private fun maybeAddIdempotencyHeader(request: HttpRequest): HttpRequest { + if (idempotencyHeader == null || request.headers.names().contains(idempotencyHeader)) { + return request + } + + return request + .toBuilder() + // Set a header to uniquely identify the request when retried. + .putHeader(idempotencyHeader, idempotencyKey()) + .build() + } + + private fun shouldRetry(response: HttpResponse): Boolean { + val headers = response.headers() + return shouldRetry(response.statusCode(), headers) + } + + private fun shouldRetry(statusCode: Int, headers: Headers): Boolean { + // Note: this is not a standard header + val shouldRetryHeader = headers.values("X-Should-Retry").getOrNull(0) + + return when { + // If the server explicitly says whether to retry, obey + shouldRetryHeader == "true" -> true + shouldRetryHeader == "false" -> false + + // Retry on request timeouts + statusCode == 408 -> true + // Retry on lock timeouts + statusCode == 409 -> true + // Retry on rate limits + statusCode == 429 -> true + // Retry internal errors + statusCode >= 500 -> true + else -> false + } + } + + private fun shouldRetry(throwable: Throwable): Boolean = + // Only retry known retryable exceptions, other exceptions are not intended to be retried. + throwable is IOException || + throwable is OpenAIIoException || + throwable is OpenAIRetryableException + + private fun getRetryBackoffDuration( + retries: Int, + response: HttpResponse? = null, + error: Throwable? = null, + ): Duration { + // About the Retry-After header: + // https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Retry-After + (response?.headers() ?: (error?.let(::unwrap) as? UnexpectedStatusCodeException)?.headers()) + ?.let { headers -> + headers + .values("Retry-After-Ms") + .getOrNull(0) + ?.toFloatOrNull() + ?.times(TimeUnit.MILLISECONDS.toNanos(1)) + ?: headers.values("Retry-After").getOrNull(0)?.let { retryAfter -> + retryAfter.toFloatOrNull()?.times(TimeUnit.SECONDS.toNanos(1)) + ?: try { + ChronoUnit.NANOS.between( + OffsetDateTime.now(clock), + OffsetDateTime.parse( + retryAfter, + DateTimeFormatter.RFC_1123_DATE_TIME, + ), + ) + } catch (e: DateTimeParseException) { + null + } + } + } + ?.let { retryAfterNanos -> + // If the API asks us to wait a certain amount of time, do what it says. + return Duration.ofNanos(retryAfterNanos.toLong()) + } + + // Apply exponential backoff, but not more than the max. + val backoffSeconds = min(0.5 * 2.0.pow(retries - 1), 8.0) + + // Apply some jitter + val jitter = 1.0 - 0.25 * ThreadLocalRandom.current().nextDouble() + + return Duration.ofNanos((TimeUnit.SECONDS.toNanos(1) * backoffSeconds * jitter).toLong()) + } +} diff --git a/openai-java-core/src/main/kotlin/com/openai/services/async/ContainerServiceAsyncImpl.kt b/openai-java-core/src/main/kotlin/com/openai/services/async/ContainerServiceAsyncImpl.kt index 614268207..9ae7c6188 100644 --- a/openai-java-core/src/main/kotlin/com/openai/services/async/ContainerServiceAsyncImpl.kt +++ b/openai-java-core/src/main/kotlin/com/openai/services/async/ContainerServiceAsyncImpl.kt @@ -15,6 +15,7 @@ import com.openai.core.http.HttpRequest import com.openai.core.http.HttpResponse import com.openai.core.http.HttpResponse.Handler import com.openai.core.http.HttpResponseFor +import com.openai.core.http.closeIfPipelineOwned import com.openai.core.http.json import com.openai.core.http.parseable import com.openai.core.prepareAsync @@ -74,7 +75,7 @@ class ContainerServiceAsyncImpl internal constructor(private val clientOptions: requestOptions: RequestOptions, ): CompletableFuture = // delete /containers/{container_id} - withRawResponse().delete(params, requestOptions).thenAccept {} + withRawResponse().delete(params, requestOptions).thenAccept { it.closeIfPipelineOwned() } class WithRawResponseImpl internal constructor(private val clientOptions: ClientOptions) : ContainerServiceAsync.WithRawResponse { diff --git a/openai-java-core/src/main/kotlin/com/openai/services/async/ResponseServiceAsyncImpl.kt b/openai-java-core/src/main/kotlin/com/openai/services/async/ResponseServiceAsyncImpl.kt index 3701198b2..fe61ee8c6 100644 --- a/openai-java-core/src/main/kotlin/com/openai/services/async/ResponseServiceAsyncImpl.kt +++ b/openai-java-core/src/main/kotlin/com/openai/services/async/ResponseServiceAsyncImpl.kt @@ -20,6 +20,7 @@ import com.openai.core.http.HttpResponse import com.openai.core.http.HttpResponse.Handler import com.openai.core.http.HttpResponseFor import com.openai.core.http.StreamResponse +import com.openai.core.http.closeIfPipelineOwned import com.openai.core.http.json import com.openai.core.http.map import com.openai.core.http.parseable @@ -105,7 +106,7 @@ class ResponseServiceAsyncImpl internal constructor(private val clientOptions: C requestOptions: RequestOptions, ): CompletableFuture = // delete /responses/{response_id} - withRawResponse().delete(params, requestOptions).thenAccept {} + withRawResponse().delete(params, requestOptions).thenAccept { it.closeIfPipelineOwned() } override fun cancel( params: ResponseCancelParams, diff --git a/openai-java-core/src/main/kotlin/com/openai/services/async/beta/ResponseServiceAsyncImpl.kt b/openai-java-core/src/main/kotlin/com/openai/services/async/beta/ResponseServiceAsyncImpl.kt index 9b9950e29..f12fb823c 100644 --- a/openai-java-core/src/main/kotlin/com/openai/services/async/beta/ResponseServiceAsyncImpl.kt +++ b/openai-java-core/src/main/kotlin/com/openai/services/async/beta/ResponseServiceAsyncImpl.kt @@ -20,6 +20,7 @@ import com.openai.core.http.HttpResponse import com.openai.core.http.HttpResponse.Handler import com.openai.core.http.HttpResponseFor import com.openai.core.http.StreamResponse +import com.openai.core.http.closeIfPipelineOwned import com.openai.core.http.json import com.openai.core.http.map import com.openai.core.http.parseable @@ -104,7 +105,7 @@ class ResponseServiceAsyncImpl internal constructor(private val clientOptions: C requestOptions: RequestOptions, ): CompletableFuture = // delete /responses/{response_id}?beta=true - withRawResponse().delete(params, requestOptions).thenAccept {} + withRawResponse().delete(params, requestOptions).thenAccept { it.closeIfPipelineOwned() } override fun cancel( params: ResponseCancelParams, diff --git a/openai-java-core/src/main/kotlin/com/openai/services/async/containers/FileServiceAsyncImpl.kt b/openai-java-core/src/main/kotlin/com/openai/services/async/containers/FileServiceAsyncImpl.kt index 6671cd08a..689b92743 100644 --- a/openai-java-core/src/main/kotlin/com/openai/services/async/containers/FileServiceAsyncImpl.kt +++ b/openai-java-core/src/main/kotlin/com/openai/services/async/containers/FileServiceAsyncImpl.kt @@ -15,6 +15,7 @@ import com.openai.core.http.HttpRequest import com.openai.core.http.HttpResponse import com.openai.core.http.HttpResponse.Handler import com.openai.core.http.HttpResponseFor +import com.openai.core.http.closeIfPipelineOwned import com.openai.core.http.json import com.openai.core.http.multipartFormData import com.openai.core.http.parseable @@ -75,7 +76,7 @@ class FileServiceAsyncImpl internal constructor(private val clientOptions: Clien requestOptions: RequestOptions, ): CompletableFuture = // delete /containers/{container_id}/files/{file_id} - withRawResponse().delete(params, requestOptions).thenAccept {} + withRawResponse().delete(params, requestOptions).thenAccept { it.closeIfPipelineOwned() } class WithRawResponseImpl internal constructor(private val clientOptions: ClientOptions) : FileServiceAsync.WithRawResponse { diff --git a/openai-java-core/src/main/kotlin/com/openai/services/async/realtime/CallServiceAsyncImpl.kt b/openai-java-core/src/main/kotlin/com/openai/services/async/realtime/CallServiceAsyncImpl.kt index c7c97b379..b8d979956 100644 --- a/openai-java-core/src/main/kotlin/com/openai/services/async/realtime/CallServiceAsyncImpl.kt +++ b/openai-java-core/src/main/kotlin/com/openai/services/async/realtime/CallServiceAsyncImpl.kt @@ -13,6 +13,7 @@ import com.openai.core.http.HttpMethod import com.openai.core.http.HttpRequest import com.openai.core.http.HttpResponse import com.openai.core.http.HttpResponse.Handler +import com.openai.core.http.closeIfPipelineOwned import com.openai.core.http.encodeMultipartFields import com.openai.core.http.json import com.openai.core.http.multipartFormData @@ -51,28 +52,28 @@ class CallServiceAsyncImpl internal constructor(private val clientOptions: Clien requestOptions: RequestOptions, ): CompletableFuture = // post /realtime/calls/{call_id}/accept - withRawResponse().accept(params, requestOptions).thenAccept {} + withRawResponse().accept(params, requestOptions).thenAccept { it.closeIfPipelineOwned() } override fun hangup( params: CallHangupParams, requestOptions: RequestOptions, ): CompletableFuture = // post /realtime/calls/{call_id}/hangup - withRawResponse().hangup(params, requestOptions).thenAccept {} + withRawResponse().hangup(params, requestOptions).thenAccept { it.closeIfPipelineOwned() } override fun refer( params: CallReferParams, requestOptions: RequestOptions, ): CompletableFuture = // post /realtime/calls/{call_id}/refer - withRawResponse().refer(params, requestOptions).thenAccept {} + withRawResponse().refer(params, requestOptions).thenAccept { it.closeIfPipelineOwned() } override fun reject( params: CallRejectParams, requestOptions: RequestOptions, ): CompletableFuture = // post /realtime/calls/{call_id}/reject - withRawResponse().reject(params, requestOptions).thenAccept {} + withRawResponse().reject(params, requestOptions).thenAccept { it.closeIfPipelineOwned() } class WithRawResponseImpl internal constructor(private val clientOptions: ClientOptions) : CallServiceAsync.WithRawResponse { diff --git a/openai-java-core/src/main/kotlin/com/openai/services/blocking/ContainerServiceImpl.kt b/openai-java-core/src/main/kotlin/com/openai/services/blocking/ContainerServiceImpl.kt index 8ac54b28c..6c76aa922 100644 --- a/openai-java-core/src/main/kotlin/com/openai/services/blocking/ContainerServiceImpl.kt +++ b/openai-java-core/src/main/kotlin/com/openai/services/blocking/ContainerServiceImpl.kt @@ -15,6 +15,7 @@ import com.openai.core.http.HttpRequest import com.openai.core.http.HttpResponse import com.openai.core.http.HttpResponse.Handler import com.openai.core.http.HttpResponseFor +import com.openai.core.http.closeIfPipelineOwned import com.openai.core.http.json import com.openai.core.http.parseable import com.openai.core.prepare @@ -70,7 +71,7 @@ class ContainerServiceImpl internal constructor(private val clientOptions: Clien override fun delete(params: ContainerDeleteParams, requestOptions: RequestOptions) { // delete /containers/{container_id} - withRawResponse().delete(params, requestOptions) + withRawResponse().delete(params, requestOptions).closeIfPipelineOwned() } class WithRawResponseImpl internal constructor(private val clientOptions: ClientOptions) : diff --git a/openai-java-core/src/main/kotlin/com/openai/services/blocking/ResponseServiceImpl.kt b/openai-java-core/src/main/kotlin/com/openai/services/blocking/ResponseServiceImpl.kt index 517f5136f..467592a34 100644 --- a/openai-java-core/src/main/kotlin/com/openai/services/blocking/ResponseServiceImpl.kt +++ b/openai-java-core/src/main/kotlin/com/openai/services/blocking/ResponseServiceImpl.kt @@ -19,6 +19,7 @@ import com.openai.core.http.HttpResponse import com.openai.core.http.HttpResponse.Handler import com.openai.core.http.HttpResponseFor import com.openai.core.http.StreamResponse +import com.openai.core.http.closeIfPipelineOwned import com.openai.core.http.json import com.openai.core.http.map import com.openai.core.http.parseable @@ -86,7 +87,7 @@ class ResponseServiceImpl internal constructor(private val clientOptions: Client override fun delete(params: ResponseDeleteParams, requestOptions: RequestOptions) { // delete /responses/{response_id} - withRawResponse().delete(params, requestOptions) + withRawResponse().delete(params, requestOptions).closeIfPipelineOwned() } override fun cancel(params: ResponseCancelParams, requestOptions: RequestOptions): Response = diff --git a/openai-java-core/src/main/kotlin/com/openai/services/blocking/beta/ResponseServiceImpl.kt b/openai-java-core/src/main/kotlin/com/openai/services/blocking/beta/ResponseServiceImpl.kt index 4c7a6f3f7..a81f8ec67 100644 --- a/openai-java-core/src/main/kotlin/com/openai/services/blocking/beta/ResponseServiceImpl.kt +++ b/openai-java-core/src/main/kotlin/com/openai/services/blocking/beta/ResponseServiceImpl.kt @@ -19,6 +19,7 @@ import com.openai.core.http.HttpResponse import com.openai.core.http.HttpResponse.Handler import com.openai.core.http.HttpResponseFor import com.openai.core.http.StreamResponse +import com.openai.core.http.closeIfPipelineOwned import com.openai.core.http.json import com.openai.core.http.map import com.openai.core.http.parseable @@ -88,7 +89,7 @@ class ResponseServiceImpl internal constructor(private val clientOptions: Client override fun delete(params: ResponseDeleteParams, requestOptions: RequestOptions) { // delete /responses/{response_id}?beta=true - withRawResponse().delete(params, requestOptions) + withRawResponse().delete(params, requestOptions).closeIfPipelineOwned() } override fun cancel( diff --git a/openai-java-core/src/main/kotlin/com/openai/services/blocking/containers/FileServiceImpl.kt b/openai-java-core/src/main/kotlin/com/openai/services/blocking/containers/FileServiceImpl.kt index 23529fa96..356d0d0f2 100644 --- a/openai-java-core/src/main/kotlin/com/openai/services/blocking/containers/FileServiceImpl.kt +++ b/openai-java-core/src/main/kotlin/com/openai/services/blocking/containers/FileServiceImpl.kt @@ -15,6 +15,7 @@ import com.openai.core.http.HttpRequest import com.openai.core.http.HttpResponse import com.openai.core.http.HttpResponse.Handler import com.openai.core.http.HttpResponseFor +import com.openai.core.http.closeIfPipelineOwned import com.openai.core.http.json import com.openai.core.http.multipartFormData import com.openai.core.http.parseable @@ -67,7 +68,7 @@ class FileServiceImpl internal constructor(private val clientOptions: ClientOpti override fun delete(params: FileDeleteParams, requestOptions: RequestOptions) { // delete /containers/{container_id}/files/{file_id} - withRawResponse().delete(params, requestOptions) + withRawResponse().delete(params, requestOptions).closeIfPipelineOwned() } class WithRawResponseImpl internal constructor(private val clientOptions: ClientOptions) : diff --git a/openai-java-core/src/main/kotlin/com/openai/services/blocking/realtime/CallServiceImpl.kt b/openai-java-core/src/main/kotlin/com/openai/services/blocking/realtime/CallServiceImpl.kt index 6dd247a17..71e5616f1 100644 --- a/openai-java-core/src/main/kotlin/com/openai/services/blocking/realtime/CallServiceImpl.kt +++ b/openai-java-core/src/main/kotlin/com/openai/services/blocking/realtime/CallServiceImpl.kt @@ -13,6 +13,7 @@ import com.openai.core.http.HttpMethod import com.openai.core.http.HttpRequest import com.openai.core.http.HttpResponse import com.openai.core.http.HttpResponse.Handler +import com.openai.core.http.closeIfPipelineOwned import com.openai.core.http.encodeMultipartFields import com.openai.core.http.json import com.openai.core.http.multipartFormData @@ -43,22 +44,22 @@ class CallServiceImpl internal constructor(private val clientOptions: ClientOpti override fun accept(params: CallAcceptParams, requestOptions: RequestOptions) { // post /realtime/calls/{call_id}/accept - withRawResponse().accept(params, requestOptions) + withRawResponse().accept(params, requestOptions).closeIfPipelineOwned() } override fun hangup(params: CallHangupParams, requestOptions: RequestOptions) { // post /realtime/calls/{call_id}/hangup - withRawResponse().hangup(params, requestOptions) + withRawResponse().hangup(params, requestOptions).closeIfPipelineOwned() } override fun refer(params: CallReferParams, requestOptions: RequestOptions) { // post /realtime/calls/{call_id}/refer - withRawResponse().refer(params, requestOptions) + withRawResponse().refer(params, requestOptions).closeIfPipelineOwned() } override fun reject(params: CallRejectParams, requestOptions: RequestOptions) { // post /realtime/calls/{call_id}/reject - withRawResponse().reject(params, requestOptions) + withRawResponse().reject(params, requestOptions).closeIfPipelineOwned() } class WithRawResponseImpl internal constructor(private val clientOptions: ClientOptions) : diff --git a/openai-java-core/src/test/kotlin/com/openai/core/CancellationPropagatingFutureTest.kt b/openai-java-core/src/test/kotlin/com/openai/core/CancellationPropagatingFutureTest.kt new file mode 100644 index 000000000..4bb1015f1 --- /dev/null +++ b/openai-java-core/src/test/kotlin/com/openai/core/CancellationPropagatingFutureTest.kt @@ -0,0 +1,713 @@ +package com.openai.core + +import com.openai.client.OpenAIClientAsyncImpl +import com.openai.client.OpenAIClientImpl +import com.openai.core.http.AuthenticatedHttpRequest +import com.openai.core.http.CloseOncePipelineHttpResponse +import com.openai.core.http.Headers +import com.openai.core.http.HttpClient +import com.openai.core.http.HttpRequest +import com.openai.core.http.HttpRequestAttemptAuthenticator +import com.openai.core.http.HttpResponse +import com.openai.errors.BadRequestException +import java.io.ByteArrayInputStream +import java.io.IOException +import java.io.InputStream +import java.time.Duration +import java.util.concurrent.CompletableFuture +import java.util.concurrent.CountDownLatch +import java.util.concurrent.ExecutionException +import java.util.concurrent.Executor +import java.util.concurrent.TimeUnit +import java.util.concurrent.atomic.AtomicBoolean +import java.util.concurrent.atomic.AtomicInteger +import java.util.concurrent.atomic.AtomicReference +import org.assertj.core.api.Assertions.assertThat +import org.junit.jupiter.api.Test +import org.junit.jupiter.api.assertThrows + +internal class CancellationPropagatingFutureTest { + @Test + fun cancellationBeforeAsyncCallbackClosesCompletedInput() { + val tasks = ArrayDeque() + val closes = AtomicInteger() + val input = + CloseOncePipelineHttpResponse( + object : HttpResponse { + override fun statusCode(): Int = 200 + + override fun headers(): Headers = Headers.builder().build() + + override fun body(): InputStream = ByteArrayInputStream(ByteArray(0)) + + override fun close() { + closes.incrementAndGet() + } + } + ) + val source = CancellationPropagatingFuture.completed(input, Executor { tasks.addLast(it) }) + val result = source.thenComposeAsync { CompletableFuture.completedFuture(Unit) } + + assertThat(tasks).hasSize(1) + assertThat(result.cancel(true)).isTrue() + tasks.removeFirst().run() + + assertThat(result.isCancelled).isTrue() + assertThat(closes.get()).isEqualTo(1) + } + + @Test + fun failedComposedCancellationObservesAndClosesConcurrentSuccess() { + val entered = CountDownLatch(1) + val release = CountDownLatch(1) + val closed = CountDownLatch(1) + val closes = AtomicInteger() + val closeable = AutoCloseable { + closes.incrementAndGet() + closed.countDown() + } + val completesOnCancel = + object : CompletableFuture() { + override fun cancel(mayInterruptIfRunning: Boolean): Boolean { + complete(closeable) + return false + } + } + val result = + CancellationPropagatingFuture.completed(Unit).thenComposeAsync { + entered.countDown() + release.await(5, TimeUnit.SECONDS) + completesOnCancel + } + + assertThat(entered.await(5, TimeUnit.SECONDS)).isTrue() + assertThat(result.cancel(true)).isTrue() + release.countDown() + + assertThat(closed.await(5, TimeUnit.SECONDS)).isTrue() + assertThat(closes.get()).isEqualTo(1) + } + + @Test + fun cancellationClosesComposedValueThatLosesCompletionRace() { + val entered = CountDownLatch(1) + val release = CountDownLatch(1) + val closed = CountDownLatch(1) + val closes = AtomicInteger() + val result = + CancellationPropagatingFuture.completed(Unit).thenComposeAsync { + entered.countDown() + release.await(5, TimeUnit.SECONDS) + CompletableFuture.completedFuture( + AutoCloseable { + closes.incrementAndGet() + closed.countDown() + } + ) + } + + assertThat(entered.await(5, TimeUnit.SECONDS)).isTrue() + result.cancel(true) + release.countDown() + + assertThat(result.isCancelled).isTrue() + assertThat(closed.await(5, TimeUnit.SECONDS)).isTrue() + assertThat(closes.get()).isEqualTo(1) + } + + @Test + fun cancellationOfVoidDependentCancelsComposedOperation() { + val entered = CountDownLatch(1) + val composed = CompletableFuture() + val result = + CancellationPropagatingFuture.completed(Unit) + .thenComposeAsync { + entered.countDown() + composed + } + .thenAccept {} + + assertThat(entered.await(5, TimeUnit.SECONDS)).isTrue() + result.cancel(true) + + assertThat(result.isCancelled).isTrue() + assertThat(composed.isCancelled).isTrue() + } + + @Test + fun cancellationClosesInputWhileMapperIsRunning() { + val entered = CountDownLatch(1) + val closed = CountDownLatch(1) + val finished = CountDownLatch(1) + val closes = AtomicInteger() + val input = + CloseOncePipelineHttpResponse( + object : HttpResponse { + override fun statusCode(): Int = 200 + + override fun headers(): Headers = Headers.builder().build() + + override fun body() = ByteArrayInputStream(ByteArray(0)) + + override fun close() { + closes.incrementAndGet() + closed.countDown() + } + } + ) + val pending = CompletableFuture() + val source = CancellationPropagatingFuture.completed(Unit).thenComposeAsync { pending } + val result = + source.thenApply { response -> + try { + entered.countDown() + closed.await(5, TimeUnit.SECONDS) + } finally { + response.close() + finished.countDown() + } + } + + CompletableFuture.runAsync { pending.complete(input) } + assertThat(entered.await(5, TimeUnit.SECONDS)).isTrue() + result.cancel(true) + + assertThat(result.isCancelled).isTrue() + assertThat(closed.await(5, TimeUnit.SECONDS)).isTrue() + assertThat(finished.await(5, TimeUnit.SECONDS)).isTrue() + assertThat(closes.get()).isEqualTo(1) + } + + @Test + fun ordinaryPublicAsyncCancellationKeepsLegacyDetachedTransport() { + val dispatched = CountDownLatch(1) + val transportResult = CompletableFuture() + val transport = + object : HttpClient { + override fun execute( + request: HttpRequest, + requestOptions: RequestOptions, + ): HttpResponse = error("sync path not expected") + + override fun executeAsync( + request: HttpRequest, + requestOptions: RequestOptions, + ): CompletableFuture { + dispatched.countDown() + return transportResult + } + + override fun close() {} + } + val client = + OpenAIClientAsyncImpl( + ClientOptions.builder() + .apiKey("test-api-key") + .baseUrl("https://example.test/v1") + .httpClient(transport) + .maxRetries(0) + .build() + ) + + try { + val result = client.files().list() + assertThat(dispatched.await(5, TimeUnit.SECONDS)).isTrue() + + assertThat(result.cancel(true)).isTrue() + + assertThat(result.isCancelled).isTrue() + assertThat(transportResult.isCancelled).isFalse() + transportResult.completeExceptionally(IOException("test cleanup")) + } finally { + client.close() + } + } + + @Test + fun x509PublicAsyncCancellationReachesPendingTransport() { + val dispatched = CountDownLatch(1) + val transportResult = CancellationObservedFuture() + val transport = + object : HttpClient { + override fun execute( + request: HttpRequest, + requestOptions: RequestOptions, + ): HttpResponse = error("sync path not expected") + + override fun executeAsync( + request: HttpRequest, + requestOptions: RequestOptions, + ): CompletableFuture { + dispatched.countDown() + return transportResult + } + + override fun close() {} + } + val client = fixedAsyncClient(transport) + + try { + val result = client.files().list() + assertThat(dispatched.await(5, TimeUnit.SECONDS)).isTrue() + + assertThat(result.cancel(true)).isTrue() + + assertThat(transportResult.cancelled.await(5, TimeUnit.SECONDS)).isTrue() + assertThat(transportResult.isCancelled).isTrue() + } finally { + client.close() + } + } + + @Test + fun x509CancellationClosesLateResponseFromCancelResistantTransportOnce() { + val dispatched = CountDownLatch(1) + val closed = CountDownLatch(1) + val closes = AtomicInteger() + val transportResult = CancelResistantFuture() + val response = + object : HttpResponse { + override fun statusCode(): Int = 200 + + override fun headers(): Headers = Headers.builder().build() + + override fun body(): InputStream = ByteArrayInputStream(ByteArray(0)) + + override fun close() { + closes.incrementAndGet() + closed.countDown() + } + } + val transport = + object : HttpClient { + override fun execute( + request: HttpRequest, + requestOptions: RequestOptions, + ): HttpResponse = error("sync path not expected") + + override fun executeAsync( + request: HttpRequest, + requestOptions: RequestOptions, + ): CompletableFuture { + dispatched.countDown() + return transportResult + } + + override fun close() {} + } + val client = fixedAsyncClient(transport) + + try { + val result = client.files().list() + assertThat(dispatched.await(5, TimeUnit.SECONDS)).isTrue() + assertThat(result.cancel(true)).isTrue() + + transportResult.complete(response) + + assertThat(closed.await(5, TimeUnit.SECONDS)).isTrue() + assertThat(closes.get()).isEqualTo(1) + } finally { + client.close() + } + assertThat(closes.get()).isEqualTo(1) + } + + @Test + fun publicAsyncClientCancellationClosesResponseAndUnblocksParsing() { + val readStarted = CountDownLatch(1) + val readFinished = CountDownLatch(1) + val releaseRead = CountDownLatch(1) + val bodyClosed = AtomicBoolean() + val closes = AtomicInteger() + val body = + object : InputStream() { + override fun read(): Int { + readStarted.countDown() + try { + releaseRead.await(30, TimeUnit.SECONDS) + if (bodyClosed.get()) throw IOException("response closed") + return -1 + } finally { + readFinished.countDown() + } + } + + override fun close() { + bodyClosed.set(true) + releaseRead.countDown() + } + } + val response = + object : HttpResponse { + override fun statusCode(): Int = 200 + + override fun headers(): Headers = + Headers.builder().put("Content-Type", "application/json").build() + + override fun body(): InputStream = body + + override fun close() { + closes.incrementAndGet() + body.close() + } + } + val transportResult = CompletableFuture() + val transport = + object : HttpClient { + override fun execute( + request: HttpRequest, + requestOptions: RequestOptions, + ): HttpResponse = error("sync path not expected") + + override fun executeAsync( + request: HttpRequest, + requestOptions: RequestOptions, + ): CompletableFuture = transportResult + + override fun close() {} + } + val authenticator = + object : HttpRequestAttemptAuthenticator { + override fun authenticate( + request: HttpRequest, + timeout: Duration?, + ): AuthenticatedHttpRequest = AuthenticatedHttpRequest.create(request) {} + } + val client = + OpenAIClientAsyncImpl( + ClientOptions.builder() + .fixedBearerAuthentication("https://example.test/v1") + .fixedBearerTransport(transport, authenticator) + .maxRetries(0) + .build() + ) + + try { + val result = client.files().list() + CompletableFuture.runAsync { transportResult.complete(response) } + assertThat(readStarted.await(5, TimeUnit.SECONDS)).isTrue() + + assertThat(result.cancel(true)).isTrue() + + assertThat(result.isCancelled).isTrue() + assertThat(readFinished.await(5, TimeUnit.SECONDS)).isTrue() + assertThat(closes.get()).isEqualTo(1) + } finally { + client.close() + } + } + + @Test + fun publicBlockingX509ErrorResponseIsClosedOnce() { + val closes = AtomicInteger() + val response = + object : HttpResponse { + override fun statusCode(): Int = 400 + + override fun headers(): Headers = + Headers.builder().put("Content-Type", "application/json").build() + + override fun body(): InputStream = + ByteArrayInputStream( + """{"error":{"message":"test error","type":"invalid_request_error"}}""" + .toByteArray() + ) + + override fun close() { + closes.incrementAndGet() + } + } + val transport = + object : HttpClient { + override fun execute( + request: HttpRequest, + requestOptions: RequestOptions, + ): HttpResponse = response + + override fun executeAsync( + request: HttpRequest, + requestOptions: RequestOptions, + ): CompletableFuture = error("async path not expected") + + override fun close() {} + } + val authenticator = + object : HttpRequestAttemptAuthenticator { + override fun authenticate( + request: HttpRequest, + timeout: Duration?, + ): AuthenticatedHttpRequest = AuthenticatedHttpRequest.create(request) {} + } + val client = + OpenAIClientImpl( + ClientOptions.builder() + .fixedBearerAuthentication("https://example.test/v1") + .fixedBearerTransport(transport, authenticator) + .maxRetries(0) + .build() + ) + + try { + assertThrows { client.files().list() } + } finally { + client.close() + } + assertThat(closes.get()).isEqualTo(1) + } + + @Test + fun publicX509ErrorResponseIsClosedWhenRawResponseMappingThrows() { + val closes = AtomicInteger() + val response = + object : HttpResponse { + override fun statusCode(): Int = 400 + + override fun headers(): Headers = + Headers.builder().put("Content-Type", "application/json").build() + + override fun body(): InputStream = + ByteArrayInputStream( + """{"error":{"message":"test error","type":"invalid_request_error"}}""" + .toByteArray() + ) + + override fun close() { + closes.incrementAndGet() + } + } + val transport = + object : HttpClient { + override fun execute( + request: HttpRequest, + requestOptions: RequestOptions, + ): HttpResponse = error("sync path not expected") + + override fun executeAsync( + request: HttpRequest, + requestOptions: RequestOptions, + ): CompletableFuture = CompletableFuture.completedFuture(response) + + override fun close() {} + } + val authenticator = + object : HttpRequestAttemptAuthenticator { + override fun authenticate( + request: HttpRequest, + timeout: Duration?, + ): AuthenticatedHttpRequest = AuthenticatedHttpRequest.create(request) {} + } + val client = + OpenAIClientAsyncImpl( + ClientOptions.builder() + .fixedBearerAuthentication("https://example.test/v1") + .fixedBearerTransport(transport, authenticator) + .maxRetries(0) + .build() + ) + + try { + val result = client.files().list() + assertThrows { result.get(5, TimeUnit.SECONDS) } + } finally { + client.close() + } + assertThat(closes.get()).isEqualTo(1) + } + + @Test + fun publicOrdinaryVoidResponseKeepsLegacyCallerOwnership() { + val closes = AtomicInteger() + val response = + object : HttpResponse { + override fun statusCode(): Int = 204 + + override fun headers(): Headers = Headers.builder().build() + + override fun body(): InputStream = ByteArrayInputStream(ByteArray(0)) + + override fun close() { + closes.incrementAndGet() + } + } + val transport = + object : HttpClient { + override fun execute( + request: HttpRequest, + requestOptions: RequestOptions, + ): HttpResponse = error("sync path not expected") + + override fun executeAsync( + request: HttpRequest, + requestOptions: RequestOptions, + ): CompletableFuture = CompletableFuture.completedFuture(response) + + override fun close() {} + } + val client = + OpenAIClientAsyncImpl( + ClientOptions.builder() + .apiKey("test-api-key") + .baseUrl("https://example.test/v1") + .httpClient(transport) + .maxRetries(0) + .build() + ) + + try { + assertThat(client.responses().delete("resp_test").get(5, TimeUnit.SECONDS)).isNull() + } finally { + client.close() + } + assertThat(closes.get()).isZero() + } + + @Test + fun publicX509VoidResponseIsClosedAfterSuccessfulTerminalAction() { + val closes = AtomicInteger() + val response = + object : HttpResponse { + override fun statusCode(): Int = 204 + + override fun headers(): Headers = Headers.builder().build() + + override fun body(): InputStream = ByteArrayInputStream(ByteArray(0)) + + override fun close() { + closes.incrementAndGet() + } + } + val transport = + object : HttpClient { + override fun execute( + request: HttpRequest, + requestOptions: RequestOptions, + ): HttpResponse = error("sync path not expected") + + override fun executeAsync( + request: HttpRequest, + requestOptions: RequestOptions, + ): CompletableFuture = CompletableFuture.completedFuture(response) + + override fun close() {} + } + val authenticator = + object : HttpRequestAttemptAuthenticator { + override fun authenticate( + request: HttpRequest, + timeout: Duration?, + ): AuthenticatedHttpRequest = AuthenticatedHttpRequest.create(request) {} + } + val client = + OpenAIClientAsyncImpl( + ClientOptions.builder() + .fixedBearerAuthentication("https://example.test/v1") + .fixedBearerTransport(transport, authenticator) + .maxRetries(0) + .build() + ) + + try { + assertThat(client.responses().delete("resp_test").get(5, TimeUnit.SECONDS)).isNull() + } finally { + client.close() + } + assertThat(closes.get()).isEqualTo(1) + } + + @Test + fun publicX509RawResponseRemainsCallerOwnedAcrossDependentContinuations() { + val closes = AtomicInteger() + val body = ByteArrayInputStream("""{"object":"list","data":[]}""".toByteArray()) + val response = + object : HttpResponse { + override fun statusCode(): Int = 200 + + override fun headers(): Headers = + Headers.builder().put("Content-Type", "application/json").build() + + override fun body(): InputStream = body + + override fun close() { + closes.incrementAndGet() + } + } + val transport = + object : HttpClient { + override fun execute( + request: HttpRequest, + requestOptions: RequestOptions, + ): HttpResponse = error("sync path not expected") + + override fun executeAsync( + request: HttpRequest, + requestOptions: RequestOptions, + ): CompletableFuture = CompletableFuture.completedFuture(response) + + override fun close() {} + } + val authenticator = + object : HttpRequestAttemptAuthenticator { + override fun authenticate( + request: HttpRequest, + timeout: Duration?, + ): AuthenticatedHttpRequest = AuthenticatedHttpRequest.create(request) {} + } + val client = + OpenAIClientAsyncImpl( + ClientOptions.builder() + .fixedBearerAuthentication("https://example.test/v1") + .fixedBearerTransport(transport, authenticator) + .maxRetries(0) + .build() + ) + + try { + val raw = client.files().withRawResponse().list() + val saved = AtomicReference() + raw.thenAccept { saved.set(it) }.get(5, TimeUnit.SECONDS) + val failed = raw.thenApply { throw IllegalStateException("dependent failed") } + val failedComposed = + raw.thenComposeAsync { + throw IllegalStateException("composed dependent failed") + } + + assertThrows { failed.get(5, TimeUnit.SECONDS) } + assertThrows { failedComposed.get(5, TimeUnit.SECONDS) } + assertThat(closes.get()).isZero() + assertThat(saved.get().body().read()).isNotEqualTo(-1) + saved.get().close() + assertThat(closes.get()).isEqualTo(1) + } finally { + client.close() + } + assertThat(closes.get()).isEqualTo(1) + } + + private fun fixedAsyncClient(transport: HttpClient): OpenAIClientAsyncImpl { + val authenticator = + object : HttpRequestAttemptAuthenticator { + override fun authenticate( + request: HttpRequest, + timeout: Duration?, + ): AuthenticatedHttpRequest = AuthenticatedHttpRequest.create(request) {} + } + return OpenAIClientAsyncImpl( + ClientOptions.builder() + .fixedBearerAuthentication("https://example.test/v1") + .fixedBearerTransport(transport, authenticator) + .maxRetries(0) + .build() + ) + } + + private class CancelResistantFuture : CompletableFuture() { + override fun cancel(mayInterruptIfRunning: Boolean): Boolean = false + } + + private class CancellationObservedFuture : CompletableFuture() { + val cancelled = CountDownLatch(1) + + override fun cancel(mayInterruptIfRunning: Boolean): Boolean = + super.cancel(mayInterruptIfRunning).also { if (it) cancelled.countDown() } + } +} diff --git a/openai-java-core/src/test/kotlin/com/openai/core/CancellationPropagatingRequestBodyTest.kt b/openai-java-core/src/test/kotlin/com/openai/core/CancellationPropagatingRequestBodyTest.kt new file mode 100644 index 000000000..d021be576 --- /dev/null +++ b/openai-java-core/src/test/kotlin/com/openai/core/CancellationPropagatingRequestBodyTest.kt @@ -0,0 +1,75 @@ +package com.openai.core + +import com.openai.core.http.HttpMethod +import com.openai.core.http.HttpRequest +import com.openai.core.http.HttpRequestBody +import com.openai.core.http.withPipelineOwnedBody +import java.io.OutputStream +import java.util.concurrent.CompletableFuture +import java.util.concurrent.Executor +import org.assertj.core.api.Assertions.assertThat +import org.assertj.core.api.Assertions.assertThatThrownBy +import org.junit.jupiter.api.Test + +internal class CancellationPropagatingRequestBodyTest { + @Test + fun cancellationBeforeAsyncComposeClosesPreparedRequestBody() { + val tasks = ArrayDeque() + val executor = Executor(tasks::addLast) + val body = CountingRequestBody() + val request = + HttpRequest.builder() + .method(HttpMethod.POST) + .baseUrl("https://mtls.api.openai.com/v1") + .body(body) + .build() + .withPipelineOwnedBody() + var composed = false + val result = + CancellationPropagatingFuture.completed(request, executor).thenComposeAsync { + composed = true + CompletableFuture.completedFuture(it) + } + + result.cancel(true) + tasks.removeFirst().run() + + assertThat(composed).isFalse() + assertThat(body.closes).isEqualTo(1) + } + + @Test + fun synchronousComposeFailureClosesPreparedRequestBody() { + val body = CountingRequestBody() + val request = + HttpRequest.builder() + .method(HttpMethod.POST) + .baseUrl("https://mtls.api.openai.com/v1") + .body(body) + .build() + .withPipelineOwnedBody() + val failure = IllegalStateException("compose failed") + val result = + CancellationPropagatingFuture.completed(request, Executor(Runnable::run)) + .thenComposeAsync { throw failure } + + assertThatThrownBy(result::join).hasCause(failure) + assertThat(body.closes).isEqualTo(1) + } + + private class CountingRequestBody : HttpRequestBody { + var closes = 0 + + override fun writeTo(outputStream: OutputStream) {} + + override fun contentType(): String? = null + + override fun contentLength(): Long = 0 + + override fun repeatable(): Boolean = true + + override fun close() { + closes++ + } + } +} diff --git a/openai-java-core/src/test/kotlin/com/openai/core/ClientOptionsTest.kt b/openai-java-core/src/test/kotlin/com/openai/core/ClientOptionsTest.kt index 81c2416b8..df0ccfab0 100644 --- a/openai-java-core/src/test/kotlin/com/openai/core/ClientOptionsTest.kt +++ b/openai-java-core/src/test/kotlin/com/openai/core/ClientOptionsTest.kt @@ -5,8 +5,10 @@ import com.openai.auth.SubjectTokenProvider import com.openai.auth.SubjectTokenType import com.openai.auth.WorkloadIdentity import com.openai.azure.credential.AzureApiKeyCredential +import com.openai.core.http.AuthenticatedHttpRequest import com.openai.core.http.HttpClient import com.openai.core.http.HttpRequest +import com.openai.core.http.HttpRequestAttemptAuthenticator import com.openai.core.http.HttpRequestAuthenticator import com.openai.credential.BearerTokenCredential import com.openai.credential.WorkloadIdentityCredential @@ -182,8 +184,11 @@ internal class ClientOptionsTest { @Test fun build_withFixedBearerAuthentication_satisfiesOnlyBearerAndSurvivesCloning() { val authenticator = - object : HttpRequestAuthenticator { - override fun authenticate(request: HttpRequest): HttpRequest = request + object : HttpRequestAttemptAuthenticator { + override fun authenticate( + request: HttpRequest, + timeout: java.time.Duration?, + ): AuthenticatedHttpRequest = AuthenticatedHttpRequest.create(request) {} } val clientOptions = ClientOptions.builder() @@ -198,6 +203,8 @@ internal class ClientOptionsTest { clientOptions.securityHeaders(SecurityOptions.builder().bearerAuth(true).build()) ) .isEqualTo(com.openai.core.http.Headers.builder().build()) + assertThat(clientOptions.securityHeaders(SecurityOptions.all())) + .isEqualTo(com.openai.core.http.Headers.builder().build()) val thrown = assertThrows { clientOptions.securityHeaders( diff --git a/openai-java-core/src/test/kotlin/com/openai/core/X509BlockingResponseLifecycleTest.kt b/openai-java-core/src/test/kotlin/com/openai/core/X509BlockingResponseLifecycleTest.kt new file mode 100644 index 000000000..bb3851cc4 --- /dev/null +++ b/openai-java-core/src/test/kotlin/com/openai/core/X509BlockingResponseLifecycleTest.kt @@ -0,0 +1,118 @@ +package com.openai.core + +import com.openai.client.OpenAIClientImpl +import com.openai.core.http.AuthenticatedHttpRequest +import com.openai.core.http.Headers +import com.openai.core.http.HttpClient +import com.openai.core.http.HttpRequest +import com.openai.core.http.HttpRequestAttemptAuthenticator +import com.openai.core.http.HttpResponse +import java.io.ByteArrayInputStream +import java.io.InputStream +import java.time.Duration +import java.util.concurrent.CompletableFuture +import java.util.concurrent.atomic.AtomicInteger +import org.assertj.core.api.Assertions.assertThat +import org.assertj.core.api.Assertions.assertThatCode +import org.assertj.core.api.Assertions.assertThatThrownBy +import org.junit.jupiter.api.Test + +internal class X509BlockingResponseLifecycleTest { + @Test + fun publicX509VoidResponseClosesAfterSuccessfulTerminalAction() { + val response = CountingResponse() + val client = client(response, x509 = true) + + try { + assertThatCode { client.responses().delete("resp_test") }.doesNotThrowAnyException() + } finally { + client.close() + } + + assertThat(response.closes).hasValue(1) + } + + @Test + fun publicOrdinaryVoidResponseKeepsLegacyCallerOwnership() { + val response = CountingResponse() + val client = client(response, x509 = false) + + try { + assertThatCode { client.responses().delete("resp_test") }.doesNotThrowAnyException() + } finally { + client.close() + } + + assertThat(response.closes).hasValue(0) + } + + @Test + fun publicX509VoidResponsePropagatesCloseFailure() { + val failure = IllegalStateException("close failed") + val response = CountingResponse(failure) + val client = client(response, x509 = true) + + try { + assertThatThrownBy { client.responses().delete("resp_test") }.isSameAs(failure) + } finally { + client.close() + } + + assertThat(response.closes).hasValue(1) + } + + private fun client(response: HttpResponse, x509: Boolean): OpenAIClientImpl { + val transport = + object : HttpClient { + override fun execute( + request: HttpRequest, + requestOptions: RequestOptions, + ): HttpResponse = response + + override fun executeAsync( + request: HttpRequest, + requestOptions: RequestOptions, + ): CompletableFuture = error("async path not expected") + + override fun close() {} + } + val options = + if (x509) { + val authenticator = + object : HttpRequestAttemptAuthenticator { + override fun authenticate( + request: HttpRequest, + timeout: Duration?, + ): AuthenticatedHttpRequest = AuthenticatedHttpRequest.create(request) {} + } + ClientOptions.builder() + .fixedBearerAuthentication("https://example.test/v1") + .fixedBearerTransport(transport, authenticator) + .maxRetries(0) + .build() + } else { + ClientOptions.builder() + .apiKey("test-api-key") + .baseUrl("https://example.test/v1") + .httpClient(transport) + .maxRetries(0) + .build() + } + return OpenAIClientImpl(options) + } + + private class CountingResponse(private val closeFailure: Throwable? = null) : HttpResponse { + val closes = AtomicInteger() + + override fun statusCode(): Int = 204 + + override fun headers(): Headers = Headers.builder().build() + + override fun body(): InputStream = ByteArrayInputStream(ByteArray(0)) + + override fun close() { + closes.incrementAndGet() + closeFailure?.let { throw it } + } + } +} diff --git a/openai-java-core/src/test/kotlin/com/openai/core/handlers/PipelineErrorHandlerTest.kt b/openai-java-core/src/test/kotlin/com/openai/core/handlers/PipelineErrorHandlerTest.kt new file mode 100644 index 000000000..7417153bb --- /dev/null +++ b/openai-java-core/src/test/kotlin/com/openai/core/handlers/PipelineErrorHandlerTest.kt @@ -0,0 +1,42 @@ +package com.openai.core.handlers + +import com.openai.core.JsonField +import com.openai.core.http.Headers +import com.openai.core.http.HttpResponse +import com.openai.core.http.asPipelineOwned +import com.openai.models.ErrorObject +import java.io.ByteArrayInputStream +import org.assertj.core.api.Assertions.assertThat +import org.assertj.core.api.Assertions.assertThatThrownBy +import org.junit.jupiter.api.Test + +internal class PipelineErrorHandlerTest { + @Test + fun statusInspectionFailureClosesPipelineResponse() { + val failure = IllegalStateException("status unavailable") + var closes = 0 + val response = + object : HttpResponse { + override fun statusCode(): Int = throw failure + + override fun headers(): Headers = Headers.builder().build() + + override fun body() = ByteArrayInputStream(ByteArray(0)) + + override fun close() { + closes++ + } + } + .asPipelineOwned() + val handler = + errorHandler( + object : HttpResponse.Handler> { + override fun handle(response: HttpResponse): JsonField = + error("unused") + } + ) + + assertThatThrownBy { handler.handle(response) }.isSameAs(failure) + assertThat(closes).isEqualTo(1) + } +} diff --git a/openai-java-core/src/test/kotlin/com/openai/core/http/AttemptAuthenticatingRetryingHttpClientFixtures.kt b/openai-java-core/src/test/kotlin/com/openai/core/http/AttemptAuthenticatingRetryingHttpClientFixtures.kt new file mode 100644 index 000000000..049f7522d --- /dev/null +++ b/openai-java-core/src/test/kotlin/com/openai/core/http/AttemptAuthenticatingRetryingHttpClientFixtures.kt @@ -0,0 +1,150 @@ +package com.openai.core.http + +import com.openai.core.RequestOptions +import com.openai.core.Sleeper +import com.openai.errors.UnexpectedStatusCodeException +import java.io.ByteArrayInputStream +import java.io.OutputStream +import java.time.Duration +import java.util.concurrent.CompletableFuture + +internal fun client( + transport: HttpClient, + authenticator: HttpRequestAttemptAuthenticator, + maxRetries: Int, +): HttpClient = + RetryingHttpClient.builder() + .httpClient(transport) + .attemptAuthenticator(authenticator) + .maxRetries(maxRetries) + .sleeper(ImmediateSleeper) + .build() + +internal fun request(repeatable: Boolean = true): HttpRequest = + HttpRequest.builder() + .method(HttpMethod.POST) + .baseUrl("https://mtls.api.openai.com/v1") + .body( + object : HttpRequestBody { + override fun writeTo(outputStream: OutputStream) {} + + override fun contentLength(): Long = 0 + + override fun contentType(): String? = null + + override fun repeatable(): Boolean = repeatable + + override fun close() {} + } + ) + .build() + +internal class CachingAuthenticator(vararg failures: Throwable) : HttpRequestAttemptAuthenticator { + private val failures = ArrayDeque(failures.toList()) + private var token: String? = null + var attempts = 0 + var exchanges = 0 + var invalidations = 0 + + override fun authenticate(request: HttpRequest, timeout: Duration?): AuthenticatedHttpRequest { + attempts++ + failures.removeFirstOrNull()?.let { throw it } + val exact = token ?: "token-${++exchanges}".also { token = it } + val authenticated = + request.toBuilder().replaceHeaders("Authorization", "Bearer $exact").build() + return AuthenticatedHttpRequest.create(authenticated) { + if (token == exact) { + invalidations++ + token = null + } + } + } +} + +internal class ScriptedClient(vararg responses: HttpResponse) : HttpClient { + private val responses = ArrayDeque(responses.toList()) + val authorization = mutableListOf() + val retryCounts = mutableListOf() + val timeouts = mutableListOf() + var calls = 0 + + override fun execute(request: HttpRequest, requestOptions: RequestOptions): HttpResponse { + calls++ + timeouts += requireNotNull(requestOptions.timeout).request() + authorization += request.headers.values("Authorization").singleOrNull() + retryCounts += request.headers.values("X-Stainless-Retry-Count").singleOrNull() + return responses.removeFirst() + } + + override fun executeAsync( + request: HttpRequest, + requestOptions: RequestOptions, + ): CompletableFuture = + CompletableFuture.completedFuture(execute(request, requestOptions)) + + override fun close() {} +} + +internal object ImmediateSleeper : Sleeper { + override fun sleep(duration: Duration) {} + + override fun sleepAsync(duration: Duration): CompletableFuture = + CompletableFuture.completedFuture(null) + + override fun close() {} +} + +internal class RecordingSleeper : Sleeper { + val delays = mutableListOf() + + override fun sleep(duration: Duration) { + delays += duration + } + + override fun sleepAsync(duration: Duration): CompletableFuture { + delays += duration + return CompletableFuture.completedFuture(null) + } + + override fun close() {} +} + +internal class SlowAuthenticator(private val delay: Duration) : HttpRequestAttemptAuthenticator { + override fun authenticate(request: HttpRequest, timeout: Duration?): AuthenticatedHttpRequest { + Thread.sleep(delay.toMillis()) + return AuthenticatedHttpRequest.create(request) {} + } + + override fun authenticateAsync( + request: HttpRequest, + timeout: Duration?, + ): CompletableFuture = + CompletableFuture.completedFuture(authenticate(request, timeout)) +} + +internal fun statusFailure( + status: Int, + vararg headers: Pair, +): UnexpectedStatusCodeException = + UnexpectedStatusCodeException.builder() + .statusCode(status) + .headers( + Headers.builder() + .apply { headers.forEach { (name, value) -> put(name, value) } } + .build() + ) + .build() + +internal fun response(status: Int, vararg headers: Pair): HttpResponse = + object : HttpResponse { + override fun statusCode(): Int = status + + override fun headers(): Headers = + Headers.builder() + .apply { headers.forEach { (name, value) -> put(name, value) } } + .build() + + override fun body() = ByteArrayInputStream(ByteArray(0)) + + override fun close() {} + } diff --git a/openai-java-core/src/test/kotlin/com/openai/core/http/AttemptAuthenticatingRetryingHttpClientTest.kt b/openai-java-core/src/test/kotlin/com/openai/core/http/AttemptAuthenticatingRetryingHttpClientTest.kt new file mode 100644 index 000000000..0567fec75 --- /dev/null +++ b/openai-java-core/src/test/kotlin/com/openai/core/http/AttemptAuthenticatingRetryingHttpClientTest.kt @@ -0,0 +1,944 @@ +package com.openai.core.http + +import com.openai.core.RequestOptions +import com.openai.core.Sleeper +import com.openai.errors.OpenAIIoException +import com.openai.errors.OpenAIRetryableException +import com.openai.errors.UnexpectedStatusCodeException +import java.io.ByteArrayInputStream +import java.io.OutputStream +import java.time.Clock +import java.time.Duration +import java.util.concurrent.CompletableFuture +import java.util.concurrent.CompletionException +import java.util.concurrent.CountDownLatch +import java.util.concurrent.ExecutionException +import java.util.concurrent.TimeUnit +import java.util.concurrent.atomic.AtomicInteger +import java.util.concurrent.atomic.AtomicReference +import org.assertj.core.api.Assertions.assertThat +import org.assertj.core.api.Assertions.assertThatThrownBy +import org.junit.jupiter.api.Test + +internal class AttemptAuthenticatingRetryingHttpClientTest { + @Test + fun replaysOneUnauthorizedResponseEvenWhenTransientRetriesAreDisabled() { + val transport = ScriptedClient(response(401), response(200)) + val authenticator = CachingAuthenticator() + val client = client(transport, authenticator, maxRetries = 0) + + client.execute(request()).use { assertThat(it.statusCode()).isEqualTo(200) } + + assertThat(authenticator.exchanges).isEqualTo(2) + assertThat(authenticator.invalidations).isEqualTo(1) + assertThat(transport.authorization).containsExactly("Bearer token-1", "Bearer token-2") + assertThat(transport.retryCounts).containsExactly("0", "0") + } + + @Test + fun transientFailuresAndUnauthorizedReplayShareOneRetryLifecycle() { + val transport = + ScriptedClient( + response(500), + response(401), + response(500), + response(500), + response(200), + ) + val authenticator = CachingAuthenticator() + val client = client(transport, authenticator, maxRetries = 2) + + client.executeAsync(request()).get(5, TimeUnit.SECONDS).use { + assertThat(it.statusCode()).isEqualTo(500) + } + + assertThat(transport.calls).isEqualTo(4) + assertThat(authenticator.exchanges).isEqualTo(2) + assertThat(authenticator.invalidations).isEqualTo(1) + assertThat(transport.retryCounts).containsExactly("0", "1", "1", "2") + } + + @Test + fun exchangeFailureConsumesTheSameRetryBudgetAsApiFailure() { + val transport = ScriptedClient(response(500), response(200)) + val authenticator = CachingAuthenticator(OpenAIRetryableException("exchange failed")) + val client = client(transport, authenticator, maxRetries = 1) + + client.execute(request()).use { assertThat(it.statusCode()).isEqualTo(500) } + + assertThat(authenticator.attempts).isEqualTo(2) + assertThat(transport.calls).isEqualTo(1) + assertThat(transport.retryCounts).containsExactly("1") + } + + @Test + fun nonRepeatableRequestInvalidatesButDoesNotReplay() { + val transport = ScriptedClient(response(401), response(200)) + val authenticator = CachingAuthenticator() + val client = client(transport, authenticator, maxRetries = 2) + + client.execute(request(repeatable = false)).use { + assertThat(it.statusCode()).isEqualTo(401) + } + + assertThat(transport.calls).isEqualTo(1) + assertThat(authenticator.invalidations).isEqualTo(1) + } + + @Test + fun nonRepeatableRequestCanRetryExchangeBeforeItsOnlyApiDispatch() { + listOf(false, true).forEach { async -> + val transport = ScriptedClient(response(200)) + val authenticator = CachingAuthenticator(OpenAIRetryableException("exchange failed")) + val client = client(transport, authenticator, maxRetries = 1) + + if (async) client.executeAsync(request(repeatable = false)).get().close() + else client.execute(request(repeatable = false)).close() + + assertThat(authenticator.attempts).isEqualTo(2) + assertThat(transport.calls).isEqualTo(1) + assertThat(transport.retryCounts).containsExactly("1") + } + } + + @Test + fun requestThatBecomesNonRepeatableStopsUnauthorizedReplay() { + listOf(false, true).forEach { async -> + var repeatabilityChecks = 0 + val statefulRequest = + HttpRequest.builder() + .method(HttpMethod.POST) + .baseUrl("https://mtls.api.openai.com/v1") + .body( + object : HttpRequestBody { + override fun writeTo(outputStream: OutputStream) {} + + override fun contentLength(): Long = 0 + + override fun contentType(): String? = null + + override fun repeatable(): Boolean = ++repeatabilityChecks == 1 + + override fun close() {} + } + ) + .build() + val transport = ScriptedClient(response(401), response(401), response(200)) + val authenticator = CachingAuthenticator() + val client = client(transport, authenticator, maxRetries = 2) + + if (async) { + client.executeAsync(statefulRequest).get().use { + assertThat(it.statusCode()).isEqualTo(401) + } + } else { + client.execute(statefulRequest).use { assertThat(it.statusCode()).isEqualTo(401) } + } + + assertThat(transport.calls).isEqualTo(2) + assertThat(repeatabilityChecks).isEqualTo(2) + assertThat(authenticator.invalidations).isEqualTo(2) + } + } + + @Test + fun authenticatorReplacementWithNonRepeatableBodyStopsUnauthorizedReplay() { + listOf(false, true).forEach { async -> + val transport = ScriptedClient(response(401), response(200)) + var invalidations = 0 + val authenticator = + object : HttpRequestAttemptAuthenticator { + fun authenticated(request: HttpRequest): AuthenticatedHttpRequest { + val wireRequest = + request + .toBuilder() + .body( + object : HttpRequestBody { + override fun writeTo(outputStream: OutputStream) {} + + override fun contentLength(): Long = 0 + + override fun contentType(): String? = null + + override fun repeatable(): Boolean = false + + override fun close() {} + } + ) + .build() + return AuthenticatedHttpRequest.create(wireRequest) { invalidations++ } + } + + override fun authenticate( + request: HttpRequest, + timeout: Duration?, + ): AuthenticatedHttpRequest = authenticated(request) + + override fun authenticateAsync( + request: HttpRequest, + timeout: Duration?, + ): CompletableFuture = + CompletableFuture.completedFuture(authenticated(request)) + } + val client = client(transport, authenticator, maxRetries = 2) + + val response = + if (async) client.executeAsync(request()).get(5, TimeUnit.SECONDS) + else client.execute(request()) + response.use { assertThat(it.statusCode()).isEqualTo(401) } + + assertThat(transport.calls).isEqualTo(1) + assertThat(invalidations).isEqualTo(1) + } + } + + @Test + fun ordinaryRetriesPreservePerAttemptRequestTimeout() { + val timeout = Duration.ofMillis(123) + listOf(false, true).forEach { async -> + val transport = ScriptedClient(response(500), response(200)) + val client = + RetryingHttpClient.builder() + .httpClient(transport) + .maxRetries(1) + .sleeper(ImmediateSleeper) + .build() + val options = RequestOptions.builder().timeout(timeout).build() + + if (async) client.executeAsync(request(), options).get().close() + else client.execute(request(), options).close() + + assertThat(transport.timeouts).containsExactly(timeout, timeout) + } + } + + @Test + fun ordinaryTransportKeepsLegacyExceptionClassification() { + listOf(false, true).forEach { async -> + val failure = statusFailure(503, "X-Should-Retry" to "true") + var calls = 0 + val transport = + object : HttpClient { + override fun execute( + request: HttpRequest, + requestOptions: RequestOptions, + ): HttpResponse { + calls++ + throw failure + } + + override fun executeAsync( + request: HttpRequest, + requestOptions: RequestOptions, + ): CompletableFuture { + calls++ + return CompletableFuture().also { + it.completeExceptionally(failure) + } + } + + override fun close() {} + } + val client = + RetryingHttpClient.builder() + .httpClient(transport) + .maxRetries(2) + .sleeper(ImmediateSleeper) + .build() + + if (async) { + assertThatThrownBy { client.executeAsync(request()).get() }.hasCause(failure) + } else { + assertThatThrownBy { client.execute(request()) }.isSameAs(failure) + } + assertThat(calls).isEqualTo(1) + } + } + + @Test + fun ordinaryAsyncTransportSynchronousFailureStillEscapes() { + val failure = OpenAIRetryableException("synchronous failure") + val transport = + object : HttpClient { + override fun execute( + request: HttpRequest, + requestOptions: RequestOptions, + ): HttpResponse = error("sync path not expected") + + override fun executeAsync( + request: HttpRequest, + requestOptions: RequestOptions, + ): CompletableFuture = throw failure + + override fun close() {} + } + val ordinary = + RetryingHttpClient.builder() + .httpClient(transport) + .maxRetries(2) + .sleeper(ImmediateSleeper) + .build() + + assertThatThrownBy { ordinary.executeAsync(request()) }.isSameAs(failure) + } + + @Test + fun ordinaryAsyncRetryFactoryFailureCompletesTheReturnedFuture() { + val failure = OpenAIRetryableException("second synchronous failure") + var calls = 0 + val transport = + object : HttpClient { + override fun execute( + request: HttpRequest, + requestOptions: RequestOptions, + ): HttpResponse = error("sync path not expected") + + override fun executeAsync( + request: HttpRequest, + requestOptions: RequestOptions, + ): CompletableFuture { + if (calls++ == 0) return CompletableFuture.completedFuture(response(500)) + throw failure + } + + override fun close() {} + } + val ordinary = + RetryingHttpClient.builder() + .httpClient(transport) + .maxRetries(2) + .sleeper(ImmediateSleeper) + .build() + + assertThatThrownBy { ordinary.executeAsync(request()).get(5, TimeUnit.SECONDS) } + .hasCause(failure) + assertThat(calls).isEqualTo(2) + } + + @Test + fun ordinaryAsyncSleeperFailureKeepsLegacyRelayExceptionShape() { + val failure = OpenAIIoException("sleep failed") + val sleeper = + object : Sleeper { + override fun sleep(duration: Duration) = error("sync path not expected") + + override fun sleepAsync(duration: Duration): CompletableFuture = + CompletableFuture().also { it.completeExceptionally(failure) } + + override fun close() {} + } + val ordinary = + RetryingHttpClient.builder() + .httpClient(ScriptedClient(response(500))) + .maxRetries(1) + .sleeper(sleeper) + .build() + + val observedCompletion = AtomicReference() + val result = + ordinary.executeAsync( + request(), + RequestOptions.builder().timeout(Duration.ofSeconds(5)).build(), + ) + result.whenComplete { _, error -> observedCompletion.set(error) } + val thrown = + try { + result.get(5, TimeUnit.SECONDS) + error("expected retry sleep to fail") + } catch (error: ExecutionException) { + error + } + + assertThat(observedCompletion.get()).isInstanceOf(CompletionException::class.java) + assertThat(observedCompletion.get().cause).isSameAs(failure) + assertThat(thrown.cause).isSameAs(failure) + } + + @Test + fun ordinaryAsyncSynchronousSleeperFailureKeepsLegacyRelayExceptionShape() { + val failure = OpenAIIoException("sleep factory failed") + val sleeper = + object : Sleeper { + override fun sleep(duration: Duration) = error("sync path not expected") + + override fun sleepAsync(duration: Duration): CompletableFuture = throw failure + + override fun close() {} + } + val ordinary = + RetryingHttpClient.builder() + .httpClient(ScriptedClient(response(500))) + .maxRetries(1) + .sleeper(sleeper) + .build() + val observedCompletion = AtomicReference() + val result = + ordinary.executeAsync( + request(), + RequestOptions.builder().timeout(Duration.ofSeconds(5)).build(), + ) + result.whenComplete { _, error -> observedCompletion.set(error) } + + val thrown = + try { + result.get(5, TimeUnit.SECONDS) + error("expected retry sleep to fail") + } catch (error: ExecutionException) { + error + } + + assertThat(observedCompletion.get()).isInstanceOf(CompletionException::class.java) + assertThat(observedCompletion.get().cause).isSameAs(failure) + assertThat(thrown.cause).isSameAs(failure) + } + + @Test + fun cancellingOrdinaryAsyncResultKeepsLegacyDetachedRetryLifecycle() { + val first = CompletableFuture() + val secondDispatched = CountDownLatch(1) + val calls = AtomicInteger() + val transport = + object : HttpClient { + override fun execute( + request: HttpRequest, + requestOptions: RequestOptions, + ): HttpResponse = error("sync path not expected") + + override fun executeAsync( + request: HttpRequest, + requestOptions: RequestOptions, + ): CompletableFuture = + if (calls.incrementAndGet() == 1) first + else { + secondDispatched.countDown() + CompletableFuture.completedFuture(response(200)) + } + + override fun close() {} + } + val ordinary = + RetryingHttpClient.builder() + .httpClient(transport) + .maxRetries(1) + .sleeper(ImmediateSleeper) + .build() + val result = ordinary.executeAsync(request()) + + assertThat(result.cancel(true)).isTrue() + assertThat(first.isCancelled).isFalse() + first.complete(response(500)) + + assertThat(secondDispatched.await(5, TimeUnit.SECONDS)).isTrue() + assertThat(calls.get()).isEqualTo(2) + assertThat(result.isCancelled).isTrue() + assertThat(first.isCancelled).isFalse() + } + + @Test + fun asyncResponseInspectionFailureCompletesTheReturnedFutureAndClosesTheResponse() { + val failure = IllegalStateException("cannot inspect status") + var closes = 0 + val brokenResponse = + object : HttpResponse { + override fun statusCode(): Int = throw failure + + override fun headers(): Headers = Headers.builder().build() + + override fun body() = ByteArrayInputStream(ByteArray(0)) + + override fun close() { + closes++ + } + } + val client = client(ScriptedClient(brokenResponse), CachingAuthenticator(), maxRetries = 1) + + assertThatThrownBy { client.executeAsync(request()).get(5, TimeUnit.SECONDS) } + .hasCause(failure) + assertThat(closes).isEqualTo(1) + } + + @Test + fun syncResponseInspectionFailureClosesTheResponse() { + val failure = IllegalStateException("cannot inspect status") + var closes = 0 + val brokenResponse = + object : HttpResponse { + override fun statusCode(): Int = throw failure + + override fun headers(): Headers = Headers.builder().build() + + override fun body() = ByteArrayInputStream(ByteArray(0)) + + override fun close() { + closes++ + } + } + val client = client(ScriptedClient(brokenResponse), CachingAuthenticator(), maxRetries = 1) + + assertThatThrownBy { client.execute(request()) }.isSameAs(failure) + assertThat(closes).isEqualTo(1) + } + + @Test + fun closePreservesLegacyOrderAndRepeatabilityAfterFailure() { + val failure = OpenAIIoException("first close failed") + var transportCloses = 0 + var sleeperCloses = 0 + val transport = + object : HttpClient { + override fun execute( + request: HttpRequest, + requestOptions: RequestOptions, + ): HttpResponse = error("not expected") + + override fun executeAsync( + request: HttpRequest, + requestOptions: RequestOptions, + ): CompletableFuture = error("not expected") + + override fun close() { + if (transportCloses++ == 0) throw failure + } + } + val sleeper = + object : Sleeper { + override fun sleep(duration: Duration) = error("not expected") + + override fun sleepAsync(duration: Duration): CompletableFuture = + error("not expected") + + override fun close() { + sleeperCloses++ + } + } + val client = RetryingHttpClient.builder().httpClient(transport).sleeper(sleeper).build() + + assertThatThrownBy { client.close() }.isSameAs(failure) + assertThat(transportCloses).isEqualTo(1) + assertThat(sleeperCloses).isZero() + + client.close() + client.close() + + assertThat(transportCloses).isEqualTo(3) + assertThat(sleeperCloses).isEqualTo(2) + } + + @Test + fun authenticatedCloseAttemptsEveryOwnedComponentAfterFailure() { + val failure = OpenAIIoException("authenticator close failed") + var authenticatorCloses = 0 + var transportCloses = 0 + var sleeperCloses = 0 + val authenticator = + object : HttpRequestAttemptAuthenticator { + override fun authenticate( + request: HttpRequest, + timeout: Duration?, + ): AuthenticatedHttpRequest = AuthenticatedHttpRequest.create(request) {} + + override fun close() { + authenticatorCloses++ + throw failure + } + } + val transport = + object : HttpClient { + override fun execute( + request: HttpRequest, + requestOptions: RequestOptions, + ): HttpResponse = error("not expected") + + override fun executeAsync( + request: HttpRequest, + requestOptions: RequestOptions, + ): CompletableFuture = error("not expected") + + override fun close() { + transportCloses++ + throw failure + } + } + val sleeper = + object : Sleeper { + override fun sleep(duration: Duration) = error("not expected") + + override fun sleepAsync(duration: Duration): CompletableFuture = + error("not expected") + + override fun close() { + sleeperCloses++ + } + } + val client = + RetryingHttpClient.builder() + .httpClient(transport) + .attemptAuthenticator(authenticator) + .sleeper(sleeper) + .build() + + assertThatThrownBy { client.close() }.isSameAs(failure) + assertThat(authenticatorCloses).isEqualTo(1) + assertThat(transportCloses).isEqualTo(1) + assertThat(sleeperCloses).isEqualTo(1) + assertThat(failure.suppressed).isEmpty() + } + + @Test + fun finiteDeadlineWorksWhenMonotonicClockIsNegative() { + val times = ArrayDeque(listOf(-100L, -90L)) + val transport = ScriptedClient(response(200)) + val authenticator = CachingAuthenticator() + val client = + RetryingHttpClientOrchestrator( + httpClient = transport, + sleeper = ImmediateSleeper, + clock = Clock.systemUTC(), + maxRetries = 0, + idempotencyHeader = null, + attemptAuthenticator = authenticator, + nanoTime = { times.removeFirst() }, + ) + val options = RequestOptions.builder().timeout(Duration.ofNanos(10)).build() + + assertThatThrownBy { client.execute(request(), options) } + .isInstanceOf(OpenAIIoException::class.java) + .hasMessageContaining("deadline") + assertThat(authenticator.attempts).isZero() + assertThat(transport.calls).isZero() + } + + @Test + fun negativeOverflowingDeadlineExpiresBeforeAuthentication() { + val transport = ScriptedClient(response(200)) + val authenticator = CachingAuthenticator() + val client = client(transport, authenticator, maxRetries = 0) + val options = RequestOptions.builder().timeout(Duration.ofSeconds(Long.MIN_VALUE)).build() + + assertThatThrownBy { client.execute(request(), options) } + .isInstanceOf(OpenAIIoException::class.java) + .hasMessageContaining("deadline") + assertThat(authenticator.attempts).isZero() + assertThat(transport.calls).isZero() + } + + @Test + fun ordinaryResponseInspectionKeepsLegacyBoundaryOrderAndRetryCount() { + listOf(false, true).forEach { async -> + val events = mutableListOf() + val failure = OpenAIIoException("headers unavailable") + var calls = 0 + val first = + object : HttpResponse { + override fun statusCode(): Int { + events += "status-1" + return 500 + } + + override fun headers(): Headers { + events += "headers-1" + throw failure + } + + override fun body() = ByteArrayInputStream(ByteArray(0)) + + override fun close() {} + } + val second = + object : HttpResponse { + override fun statusCode(): Int { + events += "status-2" + return 200 + } + + override fun headers(): Headers { + events += "headers-2" + return Headers.builder().build() + } + + override fun body() = ByteArrayInputStream(ByteArray(0)) + + override fun close() {} + } + val responses = ArrayDeque(listOf(first, second)) + val retryCounts = mutableListOf() + val transport = + object : HttpClient { + override fun execute( + request: HttpRequest, + requestOptions: RequestOptions, + ): HttpResponse { + events += "transport-${++calls}" + retryCounts += + request.headers.values("X-Stainless-Retry-Count").singleOrNull() + return responses.removeFirst() + } + + override fun executeAsync( + request: HttpRequest, + requestOptions: RequestOptions, + ): CompletableFuture = + CompletableFuture.completedFuture(execute(request, requestOptions)) + + override fun close() {} + } + val client = + RetryingHttpClient.builder() + .httpClient(transport) + .maxRetries(2) + .sleeper(ImmediateSleeper) + .build() + + if (async) { + assertThatThrownBy { client.executeAsync(request()).get(5, TimeUnit.SECONDS) } + .hasCause(failure) + assertThat(events).containsExactly("transport-1", "headers-1") + assertThat(retryCounts).containsExactly("0") + } else { + client.execute(request()).close() + assertThat(events).containsExactly("transport-1", "headers-1", "transport-2") + assertThat(retryCounts).containsExactly("0", "2") + } + } + } + + @Test + fun ordinaryRepeatabilityKeepsLegacyPerAttemptOrdering() { + listOf(false, true).forEach { async -> + val events = mutableListOf() + var repeatabilityChecks = 0 + val statefulRequest = + HttpRequest.builder() + .method(HttpMethod.POST) + .baseUrl("https://api.openai.com/v1") + .body( + object : HttpRequestBody { + override fun writeTo(outputStream: OutputStream) {} + + override fun contentLength(): Long = 0 + + override fun contentType(): String? = null + + override fun repeatable(): Boolean { + events += "repeatable-${++repeatabilityChecks}" + return repeatabilityChecks == 1 + } + + override fun close() {} + } + ) + .build() + var calls = 0 + fun retryableResponse(attempt: Int): HttpResponse = + object : HttpResponse { + override fun statusCode(): Int { + events += "status-$attempt" + return 500 + } + + override fun headers(): Headers { + events += "headers-$attempt" + return Headers.builder().build() + } + + override fun body() = ByteArrayInputStream(ByteArray(0)) + + override fun close() {} + } + val transport = + object : HttpClient { + override fun execute( + request: HttpRequest, + requestOptions: RequestOptions, + ): HttpResponse { + events += "transport-${++calls}" + return retryableResponse(calls) + } + + override fun executeAsync( + request: HttpRequest, + requestOptions: RequestOptions, + ): CompletableFuture = + CompletableFuture.completedFuture(execute(request, requestOptions)) + + override fun close() {} + } + val client = + RetryingHttpClient.builder() + .httpClient(transport) + .maxRetries(2) + .sleeper(ImmediateSleeper) + .build() + + if (async) client.executeAsync(statefulRequest).get().close() + else client.execute(statefulRequest).close() + + val expected = + if (async) + listOf( + "transport-1", + "repeatable-1", + "headers-1", + "status-1", + "headers-1", + "transport-2", + "repeatable-2", + ) + else + listOf( + "repeatable-1", + "transport-1", + "headers-1", + "status-1", + "headers-1", + "repeatable-2", + "transport-2", + ) + assertThat(events).containsExactlyElementsOf(expected) + assertThat(calls).isEqualTo(2) + } + } + + @Test + fun retryAfterCannotExtendTheRequestDeadline() { + val transport = ScriptedClient(response(503, "Retry-After-Ms" to "1000")) + val authenticator = CachingAuthenticator() + val client = client(transport, authenticator, maxRetries = 2) + val options = RequestOptions.builder().timeout(Duration.ofMillis(100)).build() + + assertThatThrownBy { client.execute(request(), options) } + .isInstanceOf(OpenAIIoException::class.java) + .hasMessageContaining("deadline") + assertThat(transport.calls).isEqualTo(1) + } + + @Test + fun cancellingAsyncRequestCancelsItsActiveAuthenticationWaiter() { + val waiter = CompletableFuture() + val authenticator = + object : HttpRequestAttemptAuthenticator { + override fun authenticate( + request: HttpRequest, + timeout: Duration?, + ): AuthenticatedHttpRequest = error("sync path not expected") + + override fun authenticateAsync( + request: HttpRequest, + timeout: Duration?, + ): CompletableFuture = waiter + } + val client = client(ScriptedClient(response(200)), authenticator, maxRetries = 2) + + val result = client.executeAsync(request()) + result.cancel(true) + + assertThat(waiter.isCancelled).isTrue() + } + + @Test + fun authenticationTimeIsDeductedBeforeSyncAndAsyncApiDispatch() { + val transport = ScriptedClient(response(200), response(200)) + val authenticator = SlowAuthenticator(Duration.ofMillis(60)) + val client = client(transport, authenticator, maxRetries = 0) + val options = RequestOptions.builder().timeout(Duration.ofMillis(300)).build() + + client.execute(request(), options).close() + client.executeAsync(request(), options).get(5, TimeUnit.SECONDS).close() + + assertThat(transport.timeouts).allSatisfy { + assertThat(it).isLessThan(Duration.ofMillis(270)) + } + } + + @Test + fun exhaustedDeadlineDoesNotBecomeAnUnlimitedApiCall() { + val transport = ScriptedClient(response(200), response(200)) + val authenticator = SlowAuthenticator(Duration.ofMillis(50)) + val client = client(transport, authenticator, maxRetries = 0) + val options = RequestOptions.builder().timeout(Duration.ofMillis(10)).build() + + assertThatThrownBy { client.execute(request(), options) } + .isInstanceOf(OpenAIIoException::class.java) + .hasMessageContaining("deadline") + assertThatThrownBy { client.executeAsync(request(), options).get(5, TimeUnit.SECONDS) } + .hasCauseInstanceOf(OpenAIIoException::class.java) + assertThat(transport.calls).isZero() + } + + @Test + fun issuerRetryOverridesAndRetryAfterApplyToExchangeFailures() { + val retryAfter = RecordingSleeper() + val retrying = + RetryingHttpClient.builder() + .httpClient(ScriptedClient(response(200))) + .attemptAuthenticator( + CachingAuthenticator( + statusFailure(400, "X-Should-Retry" to "true", "Retry-After-Ms" to "17") + ) + ) + .maxRetries(1) + .sleeper(retryAfter) + .build() + + retrying.execute(request()).close() + + assertThat(retryAfter.delays).containsExactly(Duration.ofMillis(17)) + + val doNotRetry = CachingAuthenticator(statusFailure(503, "X-Should-Retry" to "false")) + val rejecting = client(ScriptedClient(response(200)), doNotRetry, maxRetries = 1) + assertThatThrownBy { rejecting.execute(request()) } + .isInstanceOf(UnexpectedStatusCodeException::class.java) + assertThat(doNotRetry.attempts).isEqualTo(1) + } + + @Test + fun synchronousAsyncStageFailuresCompleteThroughTheRetryStateMachine() { + var authenticationAttempts = 0 + val authenticator = + object : HttpRequestAttemptAuthenticator { + override fun authenticate( + request: HttpRequest, + timeout: Duration?, + ): AuthenticatedHttpRequest = AuthenticatedHttpRequest.create(request) {} + + override fun authenticateAsync( + request: HttpRequest, + timeout: Duration?, + ): CompletableFuture { + if (authenticationAttempts++ == 0) { + throw OpenAIRetryableException("synchronous authentication failure") + } + return CompletableFuture.completedFuture( + AuthenticatedHttpRequest.create(request) {} + ) + } + } + val transport = + object : HttpClient { + var attempts = 0 + + override fun execute( + request: HttpRequest, + requestOptions: RequestOptions, + ): HttpResponse = error("sync path not expected") + + override fun executeAsync( + request: HttpRequest, + requestOptions: RequestOptions, + ): CompletableFuture { + if (attempts++ == 0) { + throw OpenAIRetryableException("synchronous transport failure") + } + return CompletableFuture.completedFuture(response(200)) + } + + override fun close() {} + } + val client = client(transport, authenticator, maxRetries = 2) + + client.executeAsync(request()).get(5, TimeUnit.SECONDS).close() + + assertThat(authenticationAttempts).isEqualTo(3) + assertThat(transport.attempts).isEqualTo(2) + } +} diff --git a/openai-java-core/src/test/kotlin/com/openai/core/http/AuthenticatedRequestBodyLifecycleTest.kt b/openai-java-core/src/test/kotlin/com/openai/core/http/AuthenticatedRequestBodyLifecycleTest.kt new file mode 100644 index 000000000..a54059f23 --- /dev/null +++ b/openai-java-core/src/test/kotlin/com/openai/core/http/AuthenticatedRequestBodyLifecycleTest.kt @@ -0,0 +1,93 @@ +package com.openai.core.http + +import com.openai.errors.OpenAIIoException +import java.io.OutputStream +import java.time.Duration +import java.util.concurrent.CompletableFuture +import org.assertj.core.api.Assertions.assertThat +import org.assertj.core.api.Assertions.assertThatThrownBy +import org.junit.jupiter.api.Test + +internal class AuthenticatedRequestBodyLifecycleTest { + @Test + fun syncAuthenticationFailureClosesRequestBodyOnce() { + val failure = IllegalStateException("authentication failed") + val body = CountingRequestBody() + val authenticator = + object : HttpRequestAttemptAuthenticator { + override fun authenticate( + request: HttpRequest, + timeout: Duration?, + ): AuthenticatedHttpRequest = throw failure + } + val client = client(ScriptedClient(response(200)), authenticator, maxRetries = 0) + + assertThatThrownBy { client.execute(request(body)) }.isSameAs(failure) + assertThat(body.closes).isEqualTo(1) + } + + @Test + fun cancellingBlockedAuthenticationClosesRequestBodyOnce() { + val body = CountingRequestBody() + val authentication = CompletableFuture() + val authenticator = + object : HttpRequestAttemptAuthenticator { + override fun authenticate( + request: HttpRequest, + timeout: Duration?, + ): AuthenticatedHttpRequest = error("sync path not expected") + + override fun authenticateAsync( + request: HttpRequest, + timeout: Duration?, + ): CompletableFuture = authentication + } + val client = client(ScriptedClient(response(200)), authenticator, maxRetries = 0) + + val result = client.executeAsync(request(body)) + result.cancel(true) + + assertThat(authentication.isCancelled).isTrue() + assertThat(body.closes).isEqualTo(1) + } + + @Test + fun closedClientRejectsAndClosesRequestBody() { + val body = CountingRequestBody() + val client = client(ScriptedClient(response(200)), CachingAuthenticator(), maxRetries = 0) + client.close() + + assertThatThrownBy { client.execute(request(body)) } + .isInstanceOf(OpenAIIoException::class.java) + .hasMessageContaining("closed") + assertThat(body.closes).isEqualTo(1) + + val asyncBody = CountingRequestBody() + assertThatThrownBy { client.executeAsync(request(asyncBody)).join() } + .hasCauseInstanceOf(OpenAIIoException::class.java) + assertThat(asyncBody.closes).isEqualTo(1) + } + + private fun request(body: HttpRequestBody): HttpRequest = + HttpRequest.builder() + .method(HttpMethod.POST) + .baseUrl("https://mtls.api.openai.com/v1") + .body(body) + .build() + + private class CountingRequestBody : HttpRequestBody { + var closes = 0 + + override fun writeTo(outputStream: OutputStream) {} + + override fun contentType(): String? = null + + override fun contentLength(): Long = 0 + + override fun repeatable(): Boolean = true + + override fun close() { + closes++ + } + } +} diff --git a/openai-java-core/src/test/kotlin/com/openai/core/http/LoggingHttpClientCancellationTest.kt b/openai-java-core/src/test/kotlin/com/openai/core/http/LoggingHttpClientCancellationTest.kt new file mode 100644 index 000000000..42cf55813 --- /dev/null +++ b/openai-java-core/src/test/kotlin/com/openai/core/http/LoggingHttpClientCancellationTest.kt @@ -0,0 +1,193 @@ +package com.openai.core.http + +import com.openai.core.LogLevel +import com.openai.core.RequestOptions +import java.io.ByteArrayInputStream +import java.io.ByteArrayOutputStream +import java.io.InputStream +import java.io.PrintStream +import java.time.Clock +import java.time.Instant +import java.time.ZoneId +import java.time.ZoneOffset +import java.util.concurrent.CompletableFuture +import java.util.concurrent.TimeUnit +import org.assertj.core.api.Assertions.assertThat +import org.assertj.core.api.Assertions.assertThatThrownBy +import org.junit.jupiter.api.AfterEach +import org.junit.jupiter.api.BeforeEach +import org.junit.jupiter.api.Test +import org.junit.jupiter.api.parallel.ResourceLock + +@ResourceLock("stderr") +internal class LoggingHttpClientCancellationTest { + private lateinit var originalErr: PrintStream + + @BeforeEach + fun beforeEach() { + originalErr = System.err + System.setErr(PrintStream(ByteArrayOutputStream())) + } + + @AfterEach + fun afterEach() { + System.setErr(originalErr) + } + + @Test + fun cancellationPropagatingModeClosesResponseWhenCompletionClockFails() { + val failure = IllegalStateException("clock failed") + var reads = 0 + val clock = + object : Clock() { + override fun getZone(): ZoneId = ZoneOffset.UTC + + override fun withZone(zone: ZoneId?): Clock = this + + override fun instant(): Instant { + if (reads++ == 0) return Instant.parse("1998-04-21T00:00:00Z") + throw failure + } + } + var closes = 0 + val client = loggingClient(response { closes++ }, LogLevel.OFF, clock) + + assertThatThrownBy { client.executeAsync(request()).get(5, TimeUnit.SECONDS) } + .hasCause(failure) + assertThat(closes).isEqualTo(1) + } + + @Test + fun syncCancellationPropagatingModeClosesResponseWhenCompletionClockFails() { + val failure = IllegalStateException("clock failed") + var reads = 0 + val clock = + object : Clock() { + override fun getZone(): ZoneId = ZoneOffset.UTC + + override fun withZone(zone: ZoneId?): Clock = this + + override fun instant(): Instant { + if (reads++ == 0) return Instant.parse("1998-04-21T00:00:00Z") + throw failure + } + } + var closes = 0 + val response = response { closes++ } + val client = + LoggingHttpClient.builder() + .httpClient( + object : HttpClient { + override fun execute( + request: HttpRequest, + requestOptions: RequestOptions, + ): HttpResponse = response + + override fun executeAsync( + request: HttpRequest, + requestOptions: RequestOptions, + ): CompletableFuture = error("async path not expected") + + override fun close() {} + } + ) + .level(LogLevel.OFF) + .clock(clock) + .propagateAsyncCancellation(true) + .build() + + assertThatThrownBy { client.execute(request()) }.isSameAs(failure) + assertThat(closes).isEqualTo(1) + } + + @Test + fun debugLoggingPreservesPipelineResponseOwnership() { + var closes = 0 + val client = + loggingClient( + response { closes++ }.asPipelineOwned(), + LogLevel.DEBUG, + Clock.fixed(Instant.EPOCH, ZoneOffset.UTC), + ) + + val logged = client.executeAsync(request()).get(5, TimeUnit.SECONDS) + + assertThat(logged).isInstanceOf(PipelineOwnedResource::class.java) + logged.closeIfPipelineOwned() + logged.closeIfPipelineOwned() + assertThat(closes).isEqualTo(1) + } + + @Test + fun debugLoggingClosesPipelineResponseWhenInitializedBodyCloseFails() { + val bodyFailure = IllegalStateException("body close failed") + var responseCloses = 0 + val response = + object : HttpResponse { + override fun statusCode(): Int = 200 + + override fun headers(): Headers = Headers.builder().build() + + override fun body(): InputStream = + object : ByteArrayInputStream(ByteArray(0)) { + override fun close(): Unit = throw bodyFailure + } + + override fun close() { + responseCloses++ + } + } + val client = + loggingClient( + response.asPipelineOwned(), + LogLevel.DEBUG, + Clock.fixed(Instant.EPOCH, ZoneOffset.UTC), + ) + val logged = client.executeAsync(request()).get(5, TimeUnit.SECONDS) + logged.body() + + assertThatThrownBy { logged.closeIfPipelineOwned() }.isSameAs(bodyFailure) + logged.closeIfPipelineOwned() + assertThat(responseCloses).isEqualTo(1) + } + + private fun loggingClient( + response: HttpResponse, + level: LogLevel, + clock: Clock, + ): LoggingHttpClient = + LoggingHttpClient.builder() + .httpClient( + object : HttpClient { + override fun execute( + request: HttpRequest, + requestOptions: RequestOptions, + ): HttpResponse = error("sync path not expected") + + override fun executeAsync( + request: HttpRequest, + requestOptions: RequestOptions, + ): CompletableFuture = CompletableFuture.completedFuture(response) + + override fun close() {} + } + ) + .level(level) + .clock(clock) + .propagateAsyncCancellation(true) + .build() + + private fun response(onClose: () -> Unit): HttpResponse = + object : HttpResponse { + override fun statusCode(): Int = 200 + + override fun headers(): Headers = Headers.builder().build() + + override fun body(): InputStream = ByteArrayInputStream(ByteArray(0)) + + override fun close() = onClose() + } + + private fun request(): HttpRequest = + HttpRequest.builder().method(HttpMethod.GET).baseUrl("https://api.example.com").build() +} diff --git a/openai-java-core/src/test/kotlin/com/openai/core/http/PipelineHttpResponseForTest.kt b/openai-java-core/src/test/kotlin/com/openai/core/http/PipelineHttpResponseForTest.kt new file mode 100644 index 000000000..caba0c934 --- /dev/null +++ b/openai-java-core/src/test/kotlin/com/openai/core/http/PipelineHttpResponseForTest.kt @@ -0,0 +1,56 @@ +package com.openai.core.http + +import java.io.InputStream +import java.util.concurrent.atomic.AtomicInteger +import org.assertj.core.api.Assertions.assertThat +import org.assertj.core.api.Assertions.catchThrowable +import org.junit.jupiter.api.Test + +internal class PipelineHttpResponseForTest { + @Test + fun pipelineParseFailureClosesResponseAndPreservesPrimaryFailure() { + val parseFailure = IllegalStateException("body unavailable") + val closeFailure = IllegalStateException("close failed") + val response = ThrowingBodyResponse(parseFailure, closeFailure) + val owned = response.asPipelineOwned() + val parseable = owned.parseable { owned.body() } + + val thrown = catchThrowable { parseable.parse() } + + assertThat(thrown).isSameAs(parseFailure) + assertThat(thrown.suppressed).containsExactly(closeFailure) + assertThat(response.closes).hasValue(1) + } + + @Test + fun ordinaryParseFailureKeepsLegacyCallerOwnership() { + val parseFailure = IllegalStateException("body unavailable") + val response = ThrowingBodyResponse(parseFailure) + val parseable = response.parseable { response.body() } + + val thrown = catchThrowable { parseable.parse() } + + assertThat(thrown).isSameAs(parseFailure) + assertThat(response.closes).hasValue(0) + } + + private class ThrowingBodyResponse( + private val bodyFailure: Throwable, + private val closeFailure: Throwable? = null, + ) : HttpResponse { + val closes = AtomicInteger() + + override fun statusCode(): Int = 200 + + override fun headers(): Headers = Headers.builder().build() + + override fun body(): InputStream { + throw bodyFailure + } + + override fun close() { + closes.incrementAndGet() + closeFailure?.let { throw it } + } + } +} diff --git a/openai-java-core/src/test/kotlin/com/openai/core/http/PipelineResponseLeaseTest.kt b/openai-java-core/src/test/kotlin/com/openai/core/http/PipelineResponseLeaseTest.kt new file mode 100644 index 000000000..3f5f66707 --- /dev/null +++ b/openai-java-core/src/test/kotlin/com/openai/core/http/PipelineResponseLeaseTest.kt @@ -0,0 +1,34 @@ +package com.openai.core.http + +import com.openai.core.CancellationPropagatingFuture +import java.io.ByteArrayInputStream +import java.io.InputStream +import java.util.concurrent.CompletionException +import org.assertj.core.api.Assertions.assertThatThrownBy +import org.junit.jupiter.api.Test + +internal class PipelineResponseLeaseTest { + @Test + fun voidTerminalPropagatesPipelineResponseCloseFailure() { + val failure = IllegalStateException("close failed") + val response = + object : HttpResponse { + override fun statusCode(): Int = 200 + + override fun headers(): Headers = Headers.builder().build() + + override fun body(): InputStream = ByteArrayInputStream(ByteArray(0)) + + override fun close(): Unit = throw failure + } + + val terminal = + CancellationPropagatingFuture.completed(response.asPipelineOwned()).thenAccept { + it.closeIfPipelineOwned() + } + + assertThatThrownBy(terminal::join) + .isInstanceOf(CompletionException::class.java) + .hasCause(failure) + } +} diff --git a/openai-java-core/src/test/kotlin/com/openai/core/http/X509AsyncStreamCancellationTest.kt b/openai-java-core/src/test/kotlin/com/openai/core/http/X509AsyncStreamCancellationTest.kt new file mode 100644 index 000000000..335903ec1 --- /dev/null +++ b/openai-java-core/src/test/kotlin/com/openai/core/http/X509AsyncStreamCancellationTest.kt @@ -0,0 +1,169 @@ +package com.openai.core.http + +import com.openai.client.OpenAIClientAsyncImpl +import com.openai.core.ClientOptions +import com.openai.core.RequestOptions +import com.openai.models.responses.ResponseCreateParams +import java.io.IOException +import java.time.Duration +import java.util.concurrent.CompletableFuture +import java.util.concurrent.CountDownLatch +import java.util.concurrent.TimeUnit +import org.assertj.core.api.Assertions.assertThat +import org.junit.jupiter.api.Test + +internal class X509AsyncStreamCancellationTest { + @Test + fun closeCancelsBlockedTokenExchange() { + val authentication = CancellationObservedFuture() + val started = CountDownLatch(1) + val authenticator = + object : HttpRequestAttemptAuthenticator { + override fun authenticate( + request: HttpRequest, + timeout: Duration?, + ): AuthenticatedHttpRequest = error("sync path not expected") + + override fun authenticateAsync( + request: HttpRequest, + timeout: Duration?, + ): CompletableFuture { + started.countDown() + return authentication + } + } + val client = fixedClient(NoDispatchClient, authenticator) + + try { + val stream = client.responses().createStreaming(params()) + assertThat(started.await(5, TimeUnit.SECONDS)).isTrue() + + stream.close() + + assertThat(authentication.cancelled.await(5, TimeUnit.SECONDS)).isTrue() + assertThat(stream.onCompleteFuture().get(5, TimeUnit.SECONDS)).isNull() + } finally { + client.close() + } + } + + @Test + fun closeCancelsBlockedApiRequest() { + val response = CancellationObservedFuture() + val dispatched = CountDownLatch(1) + val transport = + object : HttpClient { + override fun execute( + request: HttpRequest, + requestOptions: RequestOptions, + ): HttpResponse = error("sync path not expected") + + override fun executeAsync( + request: HttpRequest, + requestOptions: RequestOptions, + ): CompletableFuture { + dispatched.countDown() + return response + } + + override fun close() {} + } + val authenticator = ImmediateAuthenticator() + val client = fixedClient(transport, authenticator) + + try { + val stream = client.responses().createStreaming(params()) + assertThat(dispatched.await(5, TimeUnit.SECONDS)).isTrue() + + stream.close() + + assertThat(response.cancelled.await(5, TimeUnit.SECONDS)).isTrue() + assertThat(stream.onCompleteFuture().get(5, TimeUnit.SECONDS)).isNull() + } finally { + client.close() + } + } + + @Test + fun ordinaryStreamCloseKeepsLegacyDetachedRequest() { + val response = CancellationObservedFuture() + val dispatched = CountDownLatch(1) + val transport = + object : HttpClient { + override fun execute( + request: HttpRequest, + requestOptions: RequestOptions, + ): HttpResponse = error("sync path not expected") + + override fun executeAsync( + request: HttpRequest, + requestOptions: RequestOptions, + ): CompletableFuture { + dispatched.countDown() + return response + } + + override fun close() {} + } + val client = + OpenAIClientAsyncImpl( + ClientOptions.builder() + .apiKey("test-api-key") + .baseUrl("https://example.test/v1") + .httpClient(transport) + .maxRetries(0) + .build() + ) + + try { + val stream = client.responses().createStreaming(params()) + assertThat(dispatched.await(5, TimeUnit.SECONDS)).isTrue() + + stream.close() + + assertThat(response.isCancelled).isFalse() + assertThat(stream.onCompleteFuture().get(5, TimeUnit.SECONDS)).isNull() + response.completeExceptionally(IOException("test cleanup")) + } finally { + client.close() + } + } + + private fun fixedClient(transport: HttpClient, authenticator: HttpRequestAttemptAuthenticator) = + OpenAIClientAsyncImpl( + ClientOptions.builder() + .fixedBearerAuthentication("https://example.test/v1") + .fixedBearerTransport(transport, authenticator) + .maxRetries(0) + .build() + ) + + private fun params() = + ResponseCreateParams.builder().model("gpt-4o-mini").input("Hello").build() + + private class ImmediateAuthenticator : HttpRequestAttemptAuthenticator { + override fun authenticate( + request: HttpRequest, + timeout: Duration?, + ): AuthenticatedHttpRequest = AuthenticatedHttpRequest.create(request) {} + } + + private class CancellationObservedFuture : CompletableFuture() { + val cancelled = CountDownLatch(1) + + override fun cancel(mayInterruptIfRunning: Boolean): Boolean = + super.cancel(mayInterruptIfRunning).also { if (it) cancelled.countDown() } + } + + private object NoDispatchClient : HttpClient { + override fun execute(request: HttpRequest, requestOptions: RequestOptions): HttpResponse = + error("API request must not be dispatched") + + override fun executeAsync( + request: HttpRequest, + requestOptions: RequestOptions, + ): CompletableFuture = error("API request must not be dispatched") + + override fun close() {} + } +} diff --git a/openai-java-example/src/main/java/com/openai/example/X509WorkloadIdentityExample.java b/openai-java-example/src/main/java/com/openai/example/X509WorkloadIdentityExample.java new file mode 100644 index 000000000..856d7b1d1 --- /dev/null +++ b/openai-java-example/src/main/java/com/openai/example/X509WorkloadIdentityExample.java @@ -0,0 +1,77 @@ +package com.openai.example; + +import com.openai.auth.X509WorkloadIdentity; +import com.openai.client.OpenAIClient; +import com.openai.client.okhttp.OpenAIOkHttpClient; +import com.openai.client.okhttp.X509Transport; +import java.io.InputStream; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.security.KeyStore; +import java.util.Arrays; +import javax.net.ssl.KeyManagerFactory; +import javax.net.ssl.TrustManagerFactory; +import javax.net.ssl.X509ExtendedKeyManager; +import javax.net.ssl.X509TrustManager; + +/** Uses a PKCS#12 client identity for direct-only X.509 workload identity federation. */ +public final class X509WorkloadIdentityExample { + private X509WorkloadIdentityExample() {} + + public static void main(String[] args) throws Exception { + char[] password = requireEnv("OPENAI_X509_KEYSTORE_PASSWORD").toCharArray(); + OpenAIClient client = null; + try { + KeyStore keyStore = KeyStore.getInstance("PKCS12"); + Path keyStorePath = Paths.get(requireEnv("OPENAI_X509_KEYSTORE_PATH")); + try (InputStream input = Files.newInputStream(keyStorePath)) { + keyStore.load(input, password); + } + + KeyManagerFactory keyManagers = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm()); + keyManagers.init(keyStore, password); + X509ExtendedKeyManager keyManager = Arrays.stream(keyManagers.getKeyManagers()) + .filter(X509ExtendedKeyManager.class::isInstance) + .map(X509ExtendedKeyManager.class::cast) + .findFirst() + .orElseThrow(() -> + new IllegalStateException("The PKCS#12 store did not provide an X509ExtendedKeyManager")); + + TrustManagerFactory trustManagers = + TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm()); + trustManagers.init((KeyStore) null); + X509TrustManager trustManager = Arrays.stream(trustManagers.getTrustManagers()) + .filter(X509TrustManager.class::isInstance) + .map(X509TrustManager.class::cast) + .findFirst() + .orElseThrow(() -> new IllegalStateException("The JVM did not provide an X509TrustManager")); + + X509WorkloadIdentity identity = X509WorkloadIdentity.builder() + .identityProviderId(requireEnv("OPENAI_X509_IDENTITY_PROVIDER_ID")) + .serviceAccountId(requireEnv("OPENAI_X509_SERVICE_ACCOUNT_ID")) + .build(); + X509Transport transport = X509Transport.builder() + .keyManager(keyManager) + .certificateAlias(requireEnv("OPENAI_X509_CERTIFICATE_ALIAS")) + .trustManager(trustManager) + .build(); + + client = OpenAIOkHttpClient.x509Builder(identity, transport).build(); + client.models().list(); + } finally { + Arrays.fill(password, '\0'); + if (client != null) { + client.close(); + } + } + } + + private static String requireEnv(String name) { + String value = System.getenv(name); + if (value == null || value.isEmpty()) { + throw new IllegalStateException(name + " must be set"); + } + return value; + } +} diff --git a/openai-java-runtime-compatibility/src/main/java/com/openai/compatibility/OkHttpRuntimeProbe.java b/openai-java-runtime-compatibility/src/main/java/com/openai/compatibility/OkHttpRuntimeProbe.java index a586a6347..cfb05bc81 100644 --- a/openai-java-runtime-compatibility/src/main/java/com/openai/compatibility/OkHttpRuntimeProbe.java +++ b/openai-java-runtime-compatibility/src/main/java/com/openai/compatibility/OkHttpRuntimeProbe.java @@ -1,12 +1,50 @@ package com.openai.compatibility; +import com.openai.auth.X509WorkloadIdentity; import com.openai.client.OpenAIClient; +import com.openai.client.OpenAIClientAsync; import com.openai.client.okhttp.OpenAIOkHttpClient; +import com.openai.client.okhttp.OpenAIOkHttpClientAsync; +import com.openai.client.okhttp.X509Transport; +import java.math.BigInteger; +import java.net.Socket; +import java.net.URISyntaxException; +import java.security.Principal; +import java.security.PrivateKey; +import java.security.PublicKey; +import java.security.cert.CertificateEncodingException; +import java.security.cert.CertificateException; +import java.security.cert.X509Certificate; +import java.util.Date; +import java.util.Set; +import javax.net.ssl.SSLEngine; +import javax.net.ssl.X509ExtendedKeyManager; +import javax.net.ssl.X509TrustManager; public final class OkHttpRuntimeProbe { private OkHttpRuntimeProbe() {} public static void main(String[] args) { + assertLoadedFromJar(OpenAIOkHttpClient.class); + assertLoadedFromJar(X509WorkloadIdentity.class); + assertLoadedFromJar(X509Transport.class); + X509WorkloadIdentity identity = X509WorkloadIdentity.builder() + .identityProviderId("idp_runtime_probe") + .serviceAccountId("svc_runtime_probe") + .build(); + X509Transport transport = X509Transport.builder() + .keyManager(new ProbeKeyManager()) + .certificateAlias(ProbeKeyManager.ALIAS) + .trustManager(new ProbeTrustManager()) + .build(); + + OpenAIClient x509Client = + OpenAIOkHttpClient.x509Builder(identity, transport).build(); + x509Client.close(); + OpenAIClientAsync x509AsyncClient = + OpenAIOkHttpClientAsync.x509Builder(identity, transport).build(); + x509AsyncClient.close(); + OpenAIClient client = OpenAIOkHttpClient.builder().apiKey("runtime-probe").build(); try { @@ -19,4 +57,215 @@ public static void main(String[] args) { System.out.printf("Exercised an OkHttp SDK client on Java %s.%n", System.getProperty("java.version")); } + + private static void assertLoadedFromJar(Class type) { + try { + String location = type.getProtectionDomain() + .getCodeSource() + .getLocation() + .toURI() + .getPath(); + if (!location.endsWith(".jar")) { + throw new IllegalStateException(type.getName() + " was not loaded from an installed JAR"); + } + } catch (URISyntaxException exception) { + throw new IllegalStateException("Could not inspect installed artifact location", exception); + } + } + + /** Concrete protocol objects exercise TLS/client linkage without opening a network socket. */ + private static final class ProbeKeyManager extends X509ExtendedKeyManager { + private static final String ALIAS = "runtime-probe"; + private final X509Certificate certificate = new ProbeCertificate(); + private final PrivateKey privateKey = new ProbePrivateKey(); + + @Override + public String[] getClientAliases(String keyType, Principal[] issuers) { + return new String[] {ALIAS}; + } + + @Override + public String chooseClientAlias(String[] keyTypes, Principal[] issuers, Socket socket) { + return ALIAS; + } + + @Override + public String chooseEngineClientAlias(String[] keyTypes, Principal[] issuers, SSLEngine engine) { + return ALIAS; + } + + @Override + public String[] getServerAliases(String keyType, Principal[] issuers) { + return null; + } + + @Override + public String chooseServerAlias(String keyType, Principal[] issuers, Socket socket) { + return null; + } + + @Override + public X509Certificate[] getCertificateChain(String alias) { + return ALIAS.equals(alias) ? new X509Certificate[] {certificate} : null; + } + + @Override + public PrivateKey getPrivateKey(String alias) { + return ALIAS.equals(alias) ? privateKey : null; + } + } + + private static final class ProbeTrustManager implements X509TrustManager { + @Override + public void checkClientTrusted(X509Certificate[] chain, String authType) {} + + @Override + public void checkServerTrusted(X509Certificate[] chain, String authType) {} + + @Override + public X509Certificate[] getAcceptedIssuers() { + return new X509Certificate[0]; + } + } + + private static final class ProbePrivateKey implements PrivateKey { + @Override + public String getAlgorithm() { + return "RSA"; + } + + @Override + public String getFormat() { + return "PKCS#8"; + } + + @Override + public byte[] getEncoded() { + return new byte[] {0}; + } + } + + private static final class ProbeCertificate extends X509Certificate { + @Override + public void checkValidity() {} + + @Override + public void checkValidity(Date date) {} + + @Override + public int getVersion() { + return 3; + } + + @Override + public BigInteger getSerialNumber() { + return BigInteger.ONE; + } + + @Override + public Principal getIssuerDN() { + return () -> "CN=runtime-probe"; + } + + @Override + public Principal getSubjectDN() { + return () -> "CN=runtime-probe"; + } + + @Override + public Date getNotBefore() { + return new Date(0); + } + + @Override + public Date getNotAfter() { + return new Date(Long.MAX_VALUE); + } + + @Override + public byte[] getTBSCertificate() throws CertificateEncodingException { + return new byte[] {0}; + } + + @Override + public byte[] getSignature() { + return new byte[] {0}; + } + + @Override + public String getSigAlgName() { + return "NONEwithRSA"; + } + + @Override + public String getSigAlgOID() { + return "1.2.840.113549.1.1.1"; + } + + @Override + public byte[] getSigAlgParams() { + return null; + } + + @Override + public boolean[] getIssuerUniqueID() { + return null; + } + + @Override + public boolean[] getSubjectUniqueID() { + return null; + } + + @Override + public boolean[] getKeyUsage() { + return null; + } + + @Override + public int getBasicConstraints() { + return -1; + } + + @Override + public byte[] getEncoded() throws CertificateEncodingException { + return new byte[] {0}; + } + + @Override + public void verify(PublicKey key) throws CertificateException {} + + @Override + public void verify(PublicKey key, String provider) throws CertificateException {} + + @Override + public String toString() { + return "runtime-probe-certificate"; + } + + @Override + public PublicKey getPublicKey() { + return null; + } + + @Override + public Set getCriticalExtensionOIDs() { + return null; + } + + @Override + public Set getNonCriticalExtensionOIDs() { + return null; + } + + @Override + public byte[] getExtensionValue(String oid) { + return null; + } + + @Override + public boolean hasUnsupportedCriticalExtension() { + return false; + } + } } From 4ccbbb9800595b8dd59ded8de7054012f2e243f8 Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Tue, 25 Aug 2026 23:10:30 +0000 Subject: [PATCH 3/5] fix(auth): centralize X.509 request lifecycle --- .../client/okhttp/X509ClientIntegration.kt | 310 +++++++++++++++--- .../okhttp/OpenAIOkHttpClientX509Test.kt | 108 +++++- .../okhttp/X509AttemptAuthenticatorTest.kt | 302 ++++++++++++++++- .../okhttp/X509ClientConfigurationTest.kt | 63 ++++ .../kotlin/com/openai/core/ClientOptions.kt | 38 ++- .../http/HttpRequestAttemptAuthenticator.kt | 41 +++ ...eClosingHttpRequestAttemptAuthenticator.kt | 11 + .../http/RetryingHttpClientOrchestrator.kt | 68 +++- .../core/CancellationPropagatingFutureTest.kt | 18 +- .../com/openai/core/ClientOptionsTest.kt | 3 +- .../core/X509BlockingResponseLifecycleTest.kt | 3 +- ...ticatingRetryingHttpClientLifecycleTest.kt | 120 +++++++ ...singHttpRequestAttemptAuthenticatorTest.kt | 49 +++ .../http/X509AsyncStreamCancellationTest.kt | 3 +- 14 files changed, 1029 insertions(+), 108 deletions(-) create mode 100644 openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/X509ClientConfigurationTest.kt create mode 100644 openai-java-core/src/test/kotlin/com/openai/core/http/AttemptAuthenticatingRetryingHttpClientLifecycleTest.kt create mode 100644 openai-java-core/src/test/kotlin/com/openai/core/http/PhantomReachableClosingHttpRequestAttemptAuthenticatorTest.kt diff --git a/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/X509ClientIntegration.kt b/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/X509ClientIntegration.kt index 1e2c81e70..6f8718901 100644 --- a/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/X509ClientIntegration.kt +++ b/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/X509ClientIntegration.kt @@ -2,10 +2,12 @@ package com.openai.client.okhttp import com.openai.auth.X509WorkloadIdentity import com.openai.core.ClientOptions +import com.openai.core.RequestOptions import com.openai.core.Timeout import com.openai.core.http.AuthenticatedHttpRequest import com.openai.core.http.HttpRequest import com.openai.core.http.HttpRequestAttemptAuthenticator +import com.openai.core.http.HttpRequestAttemptTimeouts import com.openai.errors.OpenAIIoException import com.openai.errors.OpenAIRetryableException import com.openai.errors.UnexpectedStatusCodeException @@ -20,6 +22,7 @@ import java.util.concurrent.ScheduledThreadPoolExecutor import java.util.concurrent.TimeUnit import java.util.concurrent.atomic.AtomicBoolean import java.util.concurrent.atomic.AtomicLong +import java.util.concurrent.atomic.AtomicReference internal const val X509_API_BASE_URL = "https://mtls.api.openai.com/v1" @@ -27,13 +30,28 @@ internal class X509ClientConfiguration private constructor( private val identity: X509WorkloadIdentity, private val bindTransport: (Timeout) -> BoundX509Transport, + private val installTransport: + (ClientOptions.Builder, OkHttpClient, HttpRequestAttemptAuthenticator) -> ClientOptions, ) { companion object { @JvmSynthetic internal fun create( identity: X509WorkloadIdentity, bindTransport: (Timeout) -> BoundX509Transport, - ) = X509ClientConfiguration(identity, bindTransport) + ) = + X509ClientConfiguration(identity, bindTransport) { options, client, authenticator -> + options.buildWithFixedBearerTransport(client, authenticator) + } + + @JvmSynthetic + internal fun createForTest( + identity: X509WorkloadIdentity, + bindTransport: (Timeout) -> BoundX509Transport, + installTransport: + ( + ClientOptions.Builder, OkHttpClient, HttpRequestAttemptAuthenticator, + ) -> ClientOptions, + ) = X509ClientConfiguration(identity, bindTransport, installTransport) } @JvmSynthetic @@ -44,29 +62,34 @@ private constructor( @JvmSynthetic fun buildClientOptions(clientOptions: ClientOptions.Builder): ClientOptions { val transport = bindTransport(clientOptions.timeout()) - return try { - clientOptions - .fixedBearerTransport( - transport.apiClient, - X509AttemptAuthenticator(identity, transport.exchangeClient), - ) - .build() - } catch (error: Throwable) { + val authenticator = try { - transport.close() - } catch (closeError: Throwable) { - if (closeError !== error) { - error.addSuppressed(closeError) - } + X509AttemptAuthenticator(identity, transport.exchangeClient) + } catch (error: Throwable) { + closeAfterFailure(error, transport::close) + throw error } + return try { + installTransport(clientOptions, transport.apiClient, authenticator) + } catch (error: Throwable) { + closeAfterFailure(error, authenticator::close) + closeAfterFailure(error, transport.apiClient::close) throw error } } } +private fun closeAfterFailure(error: Throwable, close: () -> Unit) { + try { + close() + } catch (closeError: Throwable) { + if (closeError !== error) error.addSuppressed(closeError) + } +} + /** Owns the exchange client and installs one exact, generation-scoped bearer per API attempt. */ private class X509AttemptAuthenticator( - private val exchange: () -> CompletableFuture, + private val exchange: (RequestOptions) -> CompletableFuture, private val closeExchange: () -> Unit, private val nanoTime: () -> Long, private val beforeTokenPublication: () -> Unit, @@ -95,8 +118,16 @@ private class X509AttemptAuthenticator( private val closed = AtomicBoolean() override fun authenticate(request: HttpRequest, timeout: Duration?): AuthenticatedHttpRequest { + return authenticate(request, requestTimeouts(timeout)) + } + + override fun authenticate( + request: HttpRequest, + timeouts: HttpRequestAttemptTimeouts, + ): AuthenticatedHttpRequest { + val requestOptions = requestOptions(timeouts) validateRequest(request) - val waiter = token(timeout) + val waiter = token(requestOptions) val token = try { waiter.get() @@ -114,6 +145,14 @@ private class X509AttemptAuthenticator( request: HttpRequest, timeout: Duration?, ): CompletableFuture { + return authenticateAsync(request, requestTimeouts(timeout)) + } + + override fun authenticateAsync( + request: HttpRequest, + timeouts: HttpRequestAttemptTimeouts, + ): CompletableFuture { + val requestOptions = requestOptions(timeouts) try { validateRequest(request) } catch (error: Throwable) { @@ -121,7 +160,7 @@ private class X509AttemptAuthenticator( it.completeExceptionally(error) } } - val token = token(timeout) + val token = token(requestOptions) val result = CompletableFuture() token.whenComplete { value, error -> if (error == null) result.complete(authenticated(request, value)) @@ -131,22 +170,36 @@ private class X509AttemptAuthenticator( return result } - private fun token(timeout: Duration?): CompletableFuture { - val now = nanoTime() - val state = - synchronized(lock) { - check(!closed.get()) { "X.509 authenticator is closed" } - cached - ?.takeIf { it.isBeforeRefresh(now) && !it.isExpired(now) } - ?.let { - return CompletableFuture.completedFuture(it) - } - (refresh?.takeUnless { it.result.isDone } ?: startRefresh()).also { it.waiters++ } - } + private fun requestOptions(timeouts: HttpRequestAttemptTimeouts): RequestOptions { + if ( + timeouts.connect() == null && + timeouts.read() == null && + timeouts.write() == null && + timeouts.request() == null + ) { + return RequestOptions.none() + } + val timeout = + Timeout.builder() + .connect(timeouts.connect()) + .read(timeouts.read()) + .write(timeouts.write()) + .request(timeouts.request()) + .build() + return RequestOptions.builder().timeout(timeout).build() + } + + private fun requestTimeouts(timeout: Duration?): HttpRequestAttemptTimeouts = + HttpRequestAttemptTimeouts.create(null, null, null, timeout) + + private fun token(requestOptions: RequestOptions): CompletableFuture { + val timeout = requestOptions.timeout?.request()?.takeUnless(Duration::isZero) + val started = nanoTime() val waiter = CompletableFuture() - val detached = AtomicBoolean() - fun detach() { - if (!detached.compareAndSet(false, true)) return + val attached = AtomicReference() + + fun detach(state: Refresh) { + if (!attached.compareAndSet(state, null)) return val cancel = synchronized(lock) { state.waiters-- @@ -160,17 +213,9 @@ private class X509AttemptAuthenticator( state.result.cancel(true) } } - state.result.whenComplete { value, error -> - if (error == null) waiter.complete(value) - else waiter.completeExceptionally(unwrap(error)) - } - timeout?.let { - if (it.isZero) { - waiter.completeExceptionally(OpenAIIoException("X.509 request deadline exceeded")) - detach() - return waiter - } - val timeoutTask = + + val timeoutTask = + timeout?.let { try { beforeWaiterTimeoutSchedule() scheduler.schedule( @@ -187,20 +232,124 @@ private class X509AttemptAuthenticator( if (closed.get()) OpenAIIoException("HTTP client is closed", error) else error ) - detach() - return waiter + null + } + } + waiter.whenComplete { _, _ -> + timeoutTask?.cancel(false) + attached.get()?.let(::detach) + } + if (waiter.isDone) return waiter + + val waitedForIncompatibleRefresh = AtomicBoolean() + lateinit var acquire: () -> Unit + acquire = acquire@{ + if (waiter.isDone) return@acquire + val effectiveOptions = + try { + if (waitedForIncompatibleRefresh.get()) { + remainingRequestOptions(requestOptions, started, timeout) + } else { + requestOptions + } + } catch (error: Throwable) { + waiter.completeExceptionally(error) + return@acquire + } + val requested = ExchangeTimeouts.from(effectiveOptions) + var immediate: CachedToken? = null + var joined: Refresh? = null + var awaiting: Refresh? = null + try { + synchronized(lock) { + if (closed.get()) throw OpenAIIoException("HTTP client is closed") + if (waiter.isDone) return@synchronized + val now = nanoTime() + immediate = cached?.takeIf { it.isBeforeRefresh(now) && !it.isExpired(now) } + if (immediate == null) { + val active = refresh?.takeUnless { it.result.isDone } + if (active == null || active.canServe(requested, now)) { + val selected = active ?: startRefresh(effectiveOptions, requested) + selected.waiters++ + check(attached.compareAndSet(null, selected)) + joined = selected + } else { + awaiting = active + } + } } - waiter.whenComplete { _, _ -> timeoutTask.cancel(false) } + } catch (error: Throwable) { + waiter.completeExceptionally(error) + return@acquire + } + if (waiter.isDone) { + joined?.let(::detach) + return@acquire + } + immediate?.let { + waiter.complete(it) + return@acquire + } + joined?.let { state -> + state.result.whenComplete { value, error -> + detach(state) + if (error == null) waiter.complete(value) + else waiter.completeExceptionally(unwrap(error)) + } + return@acquire + } + val active = requireNotNull(awaiting) + active.result.whenComplete { value, error -> + when { + waiter.isDone -> {} + error == null -> waiter.complete(value) + closed.get() -> + waiter.completeExceptionally(OpenAIIoException("HTTP client is closed")) + active.result.isCancelled || isTransient(error) -> { + waitedForIncompatibleRefresh.set(true) + acquire() + } + else -> waiter.completeExceptionally(unwrap(error)) + } + } } - waiter.whenComplete { _, _ -> detach() } + acquire() return waiter } - private fun startRefresh(): Refresh { + private fun remainingRequestOptions( + requestOptions: RequestOptions, + started: Long, + timeout: Duration?, + ): RequestOptions { + if (timeout == null) return requestOptions + val elapsed = elapsedSince(started, nanoTime()) + val total = timeout.toNanos() + if (elapsed < 0 || elapsed >= total) { + throw OpenAIIoException("X.509 request deadline exceeded") + } + val adjustedTimeout = + requireNotNull(requestOptions.timeout) + .toBuilder() + .request(Duration.ofNanos(total - elapsed)) + .build() + return RequestOptions.builder() + .apply { + requestOptions.responseValidation?.let { responseValidation(it) } + timeout(adjustedTimeout) + } + .build() + } + + private fun startRefresh( + requestOptions: RequestOptions, + exchangeTimeouts: ExchangeTimeouts, + ): Refresh { val exchangeStarted = nanoTime() - val raw = exchange() + val raw = exchange(requestOptions) val result = CompletableFuture() - val state = Refresh(raw, result, cached, invalidationEpoch) + val state = + Refresh(raw, result, cached, invalidationEpoch, exchangeStarted, exchangeTimeouts) refresh = state raw.whenComplete { exchanged, rawError -> var value: CachedToken? = null @@ -283,8 +432,61 @@ private class X509AttemptAuthenticator( val result: CompletableFuture, val fallback: CachedToken?, val invalidationEpoch: Long, + val startedAt: Long, + val exchangeTimeouts: ExchangeTimeouts, var waiters: Int = 0, - ) + ) { + fun canServe(requested: ExchangeTimeouts, now: Long): Boolean = + !result.isDone && exchangeTimeouts.canServe(requested, elapsedSince(startedAt, now)) + } + + private data class ExchangeTimeouts( + val connect: Duration?, + val read: Duration?, + val write: Duration?, + val request: Duration?, + ) { + fun canServe(requested: ExchangeTimeouts, elapsed: Long): Boolean = + connect == requested.connect && + phaseCompatible(read, request, requested.read, requested.request) && + phaseCompatible(write, request, requested.write, requested.request) && + covers(remaining(request, elapsed), requested.request) + + companion object { + fun from(options: RequestOptions): ExchangeTimeouts = + options.timeout?.let { + ExchangeTimeouts(it.connect(), it.read(), it.write(), it.request()) + } ?: ExchangeTimeouts(null, null, null, null) + + private fun covers(available: Duration?, requested: Duration?): Boolean = + when { + available == null || available.isZero -> true + requested == null || requested.isZero -> false + else -> available >= requested + } + + private fun remaining(timeout: Duration?, elapsed: Long): Duration? { + if (timeout == null || timeout.isZero) return timeout + val total = + try { + timeout.toNanos() + } catch (_: ArithmeticException) { + Long.MAX_VALUE + } + return if (elapsed < 0 || elapsed >= total) Duration.ZERO + else Duration.ofNanos(total - elapsed) + } + + private fun phaseCompatible( + available: Duration?, + availableRequest: Duration?, + requested: Duration?, + requestedRequest: Duration?, + ): Boolean = + available == requested || + (available == availableRequest && requested == requestedRequest) + } + } private class CachedToken( val value: String, @@ -376,12 +578,16 @@ internal fun x509AttemptAuthenticatorForTest( beforeTokenPublication: () -> Unit = {}, beforeRefreshCleared: () -> Unit = {}, beforeWaiterTimeoutSchedule: () -> Unit = {}, + exchangeOptionsObserver: (RequestOptions) -> Unit = {}, schedulerObserver: (ScheduledThreadPoolExecutor) -> Unit = {}, exchange: () -> CompletableFuture, ): HttpRequestAttemptAuthenticator { val scheduler = tokenWaitScheduler().also(schedulerObserver) return X509AttemptAuthenticator( - exchange, + { options -> + exchangeOptionsObserver(options) + exchange() + }, closeExchange, nanoTime, beforeTokenPublication, diff --git a/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/OpenAIOkHttpClientX509Test.kt b/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/OpenAIOkHttpClientX509Test.kt index 87efc234c..4a9c7e1e4 100644 --- a/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/OpenAIOkHttpClientX509Test.kt +++ b/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/OpenAIOkHttpClientX509Test.kt @@ -2,7 +2,10 @@ package com.openai.client.okhttp import com.fasterxml.jackson.databind.ObjectMapper import com.openai.auth.X509WorkloadIdentity +import com.openai.core.RequestOptions +import com.openai.core.Timeout import com.openai.credential.BearerTokenCredential +import com.openai.errors.OpenAIIoException import com.openai.models.files.FileListParams import java.net.Proxy import java.security.cert.X509Certificate @@ -431,7 +434,12 @@ internal class OpenAIOkHttpClientX509Test { @Test fun cancellingPublicAsyncFutureCancelsBlockedExchange() { Fixture().use { fixture -> - fixture.authPeer.enqueue(MockResponse().setSocketPolicy(SocketPolicy.NO_RESPONSE)) + fixture.authPeer.enqueue( + MockResponse() + .setHeader("Content-Type", "application/json") + .setBody(TOKEN_RESPONSE) + .setBodyDelay(250, TimeUnit.MILLISECONDS) + ) val client = fixture.asyncBuilder().timeout(Duration.ofSeconds(30)).build() val cancelled = client.files().list() @@ -441,6 +449,104 @@ internal class OpenAIOkHttpClientX509Test { try { assertThat(cancelled.isCancelled).isTrue() + assertThat(fixture.apiPeer.server.takeRequest(1, TimeUnit.SECONDS)).isNull() + } finally { + client.close() + } + } + } + + @Test + fun synchronousRequestTimeoutBoundsIssuerAndPreventsApiDispatch() { + Fixture().use { fixture -> + fixture.authPeer.enqueue(MockResponse().setSocketPolicy(SocketPolicy.NO_RESPONSE)) + val client = fixture.syncBuilder().build() + val options = RequestOptions.builder().timeout(Duration.ofMillis(250)).build() + val started = System.nanoTime() + + try { + assertThatThrownBy { client.files().list(options) } + .isInstanceOf(OpenAIIoException::class.java) + .hasMessageContaining("deadline") + assertThat(Duration.ofNanos(System.nanoTime() - started)) + .isLessThan(Duration.ofSeconds(3)) + assertThat(fixture.apiPeer.server.requestCount).isZero() + } finally { + client.close() + } + } + } + + @Test + fun asynchronousRequestTimeoutBoundsIssuerAndPreventsApiDispatch() { + Fixture().use { fixture -> + fixture.authPeer.enqueue( + MockResponse() + .setHeader("Content-Type", "application/json") + .setBody(TOKEN_RESPONSE) + .setBodyDelay(5, TimeUnit.SECONDS) + ) + val client = fixture.asyncBuilder().build() + val options = RequestOptions.builder().timeout(Duration.ofMillis(250)).build() + val started = System.nanoTime() + + try { + assertThatThrownBy { client.files().list(options).get(3, TimeUnit.SECONDS) } + .isInstanceOf(ExecutionException::class.java) + .hasCauseInstanceOf(OpenAIIoException::class.java) + assertThat(Duration.ofNanos(System.nanoTime() - started)) + .isLessThan(Duration.ofSeconds(3)) + assertThat(fixture.apiPeer.server.requestCount).isZero() + } finally { + client.close() + } + } + } + + @Test + fun synchronousReadTimeoutReachesIssuerThroughProductionWrapper() { + Fixture().use { fixture -> + fixture.authPeer.enqueue(MockResponse().setSocketPolicy(SocketPolicy.NO_RESPONSE)) + val client = fixture.syncBuilder().maxRetries(0).build() + val timeout = + Timeout.builder() + .read(Duration.ofMillis(250)) + .request(Duration.ofSeconds(10)) + .build() + val options = RequestOptions.builder().timeout(timeout).build() + val started = System.nanoTime() + + try { + assertThatThrownBy { client.files().list(options) } + .isInstanceOf(OpenAIIoException::class.java) + assertThat(Duration.ofNanos(System.nanoTime() - started)) + .isLessThan(Duration.ofSeconds(5)) + assertThat(fixture.apiPeer.server.requestCount).isZero() + } finally { + client.close() + } + } + } + + @Test + fun asynchronousReadTimeoutReachesIssuerThroughProductionWrapper() { + Fixture().use { fixture -> + fixture.authPeer.enqueue(MockResponse().setSocketPolicy(SocketPolicy.NO_RESPONSE)) + val client = fixture.asyncBuilder().maxRetries(0).build() + val timeout = + Timeout.builder() + .read(Duration.ofMillis(250)) + .request(Duration.ofSeconds(10)) + .build() + val options = RequestOptions.builder().timeout(timeout).build() + val started = System.nanoTime() + + try { + assertThatThrownBy { client.files().list(options).get(5, TimeUnit.SECONDS) } + .isInstanceOf(ExecutionException::class.java) + .hasCauseInstanceOf(OpenAIIoException::class.java) + assertThat(Duration.ofNanos(System.nanoTime() - started)) + .isLessThan(Duration.ofSeconds(5)) assertThat(fixture.apiPeer.server.requestCount).isZero() } finally { client.close() diff --git a/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/X509AttemptAuthenticatorTest.kt b/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/X509AttemptAuthenticatorTest.kt index 3b0b66692..c4f4b5081 100644 --- a/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/X509AttemptAuthenticatorTest.kt +++ b/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/X509AttemptAuthenticatorTest.kt @@ -1,12 +1,16 @@ package com.openai.client.okhttp import com.openai.core.RequestOptions +import com.openai.core.Timeout +import com.openai.core.http.AuthenticatedHttpRequest import com.openai.core.http.Headers import com.openai.core.http.HttpClient import com.openai.core.http.HttpMethod import com.openai.core.http.HttpRequest +import com.openai.core.http.HttpRequestAttemptTimeouts import com.openai.core.http.HttpResponse import com.openai.core.http.RetryingHttpClient +import com.openai.errors.OpenAIIoException import com.openai.errors.OpenAIRetryableException import java.io.ByteArrayInputStream import java.time.Duration @@ -18,17 +22,51 @@ import java.util.concurrent.ScheduledThreadPoolExecutor import java.util.concurrent.TimeUnit import java.util.concurrent.atomic.AtomicInteger import java.util.concurrent.atomic.AtomicLong +import java.util.concurrent.atomic.AtomicReference import org.assertj.core.api.Assertions.assertThat import org.assertj.core.api.Assertions.assertThatThrownBy import org.junit.jupiter.api.Test internal class X509AttemptAuthenticatorTest { + @Test + fun forwardsEffectiveRequestOptionsToIssuerExchange() { + val observed = CompletableFuture() + val authenticator = + x509AttemptAuthenticatorForTest(exchangeOptionsObserver = { observed.complete(it) }) { + CompletableFuture.completedFuture( + X509AccessToken("optionstoken", Duration.ofMinutes(1)) + ) + } + val timeout = + Timeout.builder() + .connect(Duration.ofMillis(11)) + .read(Duration.ofMillis(22)) + .write(Duration.ofMillis(33)) + .request(Duration.ofMillis(250)) + .build() + val timeouts = + HttpRequestAttemptTimeouts.create( + timeout.connect(), + timeout.read(), + timeout.write(), + timeout.request(), + ) + + val authenticated = authenticator.authenticateAsync(request(), timeouts) + + assertThat(authorization(authenticated.get(5, TimeUnit.SECONDS))) + .isEqualTo("Bearer optionstoken") + val issuerOptions = observed.get(5, TimeUnit.SECONDS) + assertThat(issuerOptions.timeout).isEqualTo(timeout) + authenticator.close() + } + @Test fun oneCancelledWaiterDoesNotCancelSharedExchange() { val exchange = CancellationFuture() val authenticator = x509AttemptAuthenticatorForTest { exchange } - val cancelled = authenticator.authenticateAsync(request(), Duration.ofSeconds(5)) - val surviving = authenticator.authenticateAsync(request(), Duration.ofSeconds(5)) + val cancelled = authenticator.authenticateAsync(request(), null) + val surviving = authenticator.authenticateAsync(request(), null) cancelled.cancel(true) exchange.complete(X509AccessToken("survivingtoken", Duration.ofMinutes(1))) @@ -75,6 +113,261 @@ internal class X509AttemptAuthenticatorTest { authenticator.close() } + @Test + fun shortFirstAsyncWaiterCannotTruncateLongerRefreshGeneration() { + val shortExchange = CancellationFuture() + val longExchange = CancellationFuture() + val exchanges = ArrayDeque(listOf(shortExchange, longExchange)) + val authenticator = x509AttemptAuthenticatorForTest { exchanges.removeFirst() } + + val short = authenticator.authenticateAsync(request(), Duration.ofMillis(200)) + val long = authenticator.authenticateAsync(request(), Duration.ofSeconds(5)) + + assertThatThrownBy { short.get(5, TimeUnit.SECONDS) } + .isInstanceOf(ExecutionException::class.java) + .hasMessageContaining("deadline") + assertThat(shortExchange.cancelled.await(5, TimeUnit.SECONDS)).isTrue() + longExchange.complete(X509AccessToken("longtoken", Duration.ofMinutes(1))) + assertThat(authorization(long.get(5, TimeUnit.SECONDS))).isEqualTo("Bearer longtoken") + authenticator.close() + } + + @Test + fun delayedAsyncWaiterUsesActiveRefreshRemainingBudget() { + val now = AtomicLong() + val firstExchange = CompletableFuture() + val secondExchange = CompletableFuture() + val exchanges = ArrayDeque(listOf(firstExchange, secondExchange)) + val exchangeOptions = mutableListOf() + val authenticator = + x509AttemptAuthenticatorForTest( + nanoTime = now::get, + exchangeOptionsObserver = exchangeOptions::add, + ) { + exchanges.removeFirst() + } + val first = authenticator.authenticateAsync(request(), Duration.ofSeconds(5)) + now.set(Duration.ofSeconds(3).toNanos()) + val delayed = authenticator.authenticateAsync(request(), Duration.ofSeconds(4)) + + now.set(Duration.ofSeconds(4).toNanos()) + firstExchange.completeExceptionally(OpenAIIoException("first exchange timed out")) + assertThatThrownBy { first.get(5, TimeUnit.SECONDS) } + .isInstanceOf(ExecutionException::class.java) + .hasCauseInstanceOf(OpenAIIoException::class.java) + secondExchange.complete(X509AccessToken("delayedtoken", Duration.ofMinutes(1))) + + assertThat(authorization(delayed.get(5, TimeUnit.SECONDS))).isEqualTo("Bearer delayedtoken") + assertThat(exchangeOptions.map { it.timeout?.request() }) + .containsExactly(Duration.ofSeconds(5), Duration.ofSeconds(3)) + assertThat(exchanges).isEmpty() + authenticator.close() + } + + @Test + fun exhaustedDelayedAsyncWaiterDoesNotStartAnotherRefresh() { + val now = AtomicLong() + val firstExchange = CompletableFuture() + val exchangeCalls = AtomicInteger() + val authenticator = + x509AttemptAuthenticatorForTest( + nanoTime = now::get, + exchangeOptionsObserver = { exchangeCalls.incrementAndGet() }, + ) { + firstExchange + } + val first = authenticator.authenticateAsync(request(), Duration.ofSeconds(5)) + now.set(Duration.ofSeconds(3).toNanos()) + val exhausted = authenticator.authenticateAsync(request(), Duration.ofSeconds(4)) + + now.set(Duration.ofSeconds(7).toNanos()) + firstExchange.completeExceptionally(OpenAIIoException("first exchange timed out")) + + assertThatThrownBy { first.get(5, TimeUnit.SECONDS) } + .isInstanceOf(ExecutionException::class.java) + .hasCauseInstanceOf(OpenAIIoException::class.java) + assertThatThrownBy { exhausted.get(5, TimeUnit.SECONDS) } + .isInstanceOf(ExecutionException::class.java) + .hasCauseInstanceOf(OpenAIIoException::class.java) + .hasMessageContaining("deadline") + assertThat(exchangeCalls).hasValue(1) + authenticator.close() + } + + @Test + fun shortFirstSyncWaiterCannotTruncateLongerRefreshGeneration() { + val shortExchange = CancellationFuture() + val longExchange = CancellationFuture() + val firstStarted = CountDownLatch(1) + val secondStarted = CountDownLatch(1) + val exchanges = ArrayDeque(listOf(shortExchange, longExchange)) + val exchangeCount = AtomicInteger() + val authenticator = x509AttemptAuthenticatorForTest { + if (exchangeCount.getAndIncrement() == 0) firstStarted.countDown() + else secondStarted.countDown() + exchanges.removeFirst() + } + val executor = Executors.newFixedThreadPool(2) + + try { + val short = + executor.submit { authenticator.authenticate(request(), Duration.ofMillis(200)) } + assertThat(firstStarted.await(5, TimeUnit.SECONDS)).isTrue() + val long: java.util.concurrent.Future = + executor.submit( + java.util.concurrent.Callable { + authenticator.authenticate(request(), Duration.ofSeconds(5)) + } + ) + assertThat(secondStarted.await(5, TimeUnit.SECONDS)).isTrue() + + assertThatThrownBy { short.get(5, TimeUnit.SECONDS) } + .isInstanceOf(ExecutionException::class.java) + .hasCauseInstanceOf(OpenAIIoException::class.java) + .hasMessageContaining("deadline") + assertThat(shortExchange.cancelled.await(5, TimeUnit.SECONDS)).isTrue() + longExchange.complete(X509AccessToken("longtoken", Duration.ofMinutes(1))) + assertThat(authorization(long.get(5, TimeUnit.SECONDS))).isEqualTo("Bearer longtoken") + } finally { + authenticator.close() + executor.shutdownNow() + } + } + + @Test + fun delayedSyncWaiterUsesActiveRefreshRemainingBudget() { + val now = AtomicLong() + val firstExchange = CompletableFuture() + val secondExchange = CompletableFuture() + val exchanges = ArrayDeque(listOf(firstExchange, secondExchange)) + val exchangeOptions = mutableListOf() + val firstStarted = CountDownLatch(1) + val incompatibleRefreshInspected = CountDownLatch(1) + val delayedThread = AtomicReference() + val delayedClockReads = AtomicInteger() + val authenticator = + x509AttemptAuthenticatorForTest( + nanoTime = { + if ( + Thread.currentThread() === delayedThread.get() && + delayedClockReads.incrementAndGet() == 2 + ) { + incompatibleRefreshInspected.countDown() + } + now.get() + }, + exchangeOptionsObserver = exchangeOptions::add, + ) { + firstStarted.countDown() + exchanges.removeFirst() + } + val executor = Executors.newFixedThreadPool(2) + + try { + val first = + executor.submit( + java.util.concurrent.Callable { + authenticator.authenticate(request(), Duration.ofSeconds(5)) + } + ) + assertThat(firstStarted.await(5, TimeUnit.SECONDS)).isTrue() + now.set(Duration.ofSeconds(3).toNanos()) + val delayed = + executor.submit( + java.util.concurrent.Callable { + delayedThread.set(Thread.currentThread()) + authenticator.authenticate(request(), Duration.ofSeconds(4)) + } + ) + assertThat(incompatibleRefreshInspected.await(5, TimeUnit.SECONDS)).isTrue() + + now.set(Duration.ofSeconds(4).toNanos()) + firstExchange.completeExceptionally(OpenAIIoException("first exchange timed out")) + assertThatThrownBy { first.get(5, TimeUnit.SECONDS) } + .isInstanceOf(ExecutionException::class.java) + .hasCauseInstanceOf(OpenAIIoException::class.java) + secondExchange.complete(X509AccessToken("delayedtoken", Duration.ofMinutes(1))) + + assertThat(authorization(delayed.get(5, TimeUnit.SECONDS))) + .isEqualTo("Bearer delayedtoken") + assertThat(exchangeOptions.map { it.timeout?.request() }) + .containsExactly(Duration.ofSeconds(5), Duration.ofSeconds(3)) + assertThat(exchanges).isEmpty() + } finally { + authenticator.close() + executor.shutdownNow() + } + } + + @Test + fun closeFailsAttachedAndIncompatibleAsyncWaitersAsIo() { + val exchange = CancellationFuture() + val exchanges = AtomicInteger() + val authenticator = x509AttemptAuthenticatorForTest { + exchanges.incrementAndGet() + exchange + } + val attached = authenticator.authenticateAsync(request(), Duration.ofSeconds(1)) + val awaiting = authenticator.authenticateAsync(request(), Duration.ofSeconds(5)) + + authenticator.close() + + listOf(attached, awaiting).forEach { waiter -> + assertThatThrownBy { waiter.get(5, TimeUnit.SECONDS) } + .isInstanceOf(ExecutionException::class.java) + .hasCauseInstanceOf(OpenAIIoException::class.java) + .hasMessageContaining("closed") + } + assertThat(exchanges).hasValue(1) + } + + @Test + fun closeFailsAttachedAndIncompatibleSyncWaitersAsIo() { + val exchange = CancellationFuture() + val scheduled = CountDownLatch(2) + val exchangeStarted = CountDownLatch(1) + val exchanges = AtomicInteger() + val authenticator = + x509AttemptAuthenticatorForTest( + beforeWaiterTimeoutSchedule = { scheduled.countDown() } + ) { + exchanges.incrementAndGet() + exchangeStarted.countDown() + exchange + } + val executor = Executors.newFixedThreadPool(2) + + try { + val attached = + executor.submit( + java.util.concurrent.Callable { + authenticator.authenticate(request(), Duration.ofSeconds(1)) + } + ) + assertThat(exchangeStarted.await(5, TimeUnit.SECONDS)).isTrue() + val awaiting = + executor.submit( + java.util.concurrent.Callable { + authenticator.authenticate(request(), Duration.ofSeconds(5)) + } + ) + assertThat(scheduled.await(5, TimeUnit.SECONDS)).isTrue() + + authenticator.close() + + listOf(attached, awaiting).forEach { waiter -> + assertThatThrownBy { waiter.get(5, TimeUnit.SECONDS) } + .isInstanceOf(ExecutionException::class.java) + .hasCauseInstanceOf(OpenAIIoException::class.java) + .hasMessageContaining("closed") + } + assertThat(exchanges).hasValue(1) + } finally { + authenticator.close() + executor.shutdownNow() + } + } + @Test fun interruptingSyncWaiterCancelsItsLastExchange() { val exchange = CancellationFuture() @@ -188,11 +481,13 @@ internal class X509AttemptAuthenticatorTest { @Test fun closeBeforeTimeoutSchedulingDoesNotExposeSchedulerRejection() { val exchange = CancellationFuture() + val exchangeStarts = AtomicInteger() lateinit var authenticator: com.openai.core.http.HttpRequestAttemptAuthenticator authenticator = x509AttemptAuthenticatorForTest( beforeWaiterTimeoutSchedule = { authenticator.close() } ) { + exchangeStarts.incrementAndGet() exchange } @@ -201,7 +496,8 @@ internal class X509AttemptAuthenticatorTest { assertThatThrownBy { authentication.get(5, TimeUnit.SECONDS) } .isInstanceOf(ExecutionException::class.java) .hasMessageContaining("HTTP client is closed") - assertThat(exchange.isCancelled).isTrue() + assertThat(exchangeStarts).hasValue(0) + assertThat(exchange.isCancelled).isFalse() } @Test diff --git a/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/X509ClientConfigurationTest.kt b/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/X509ClientConfigurationTest.kt new file mode 100644 index 000000000..254e56d2b --- /dev/null +++ b/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/X509ClientConfigurationTest.kt @@ -0,0 +1,63 @@ +package com.openai.client.okhttp + +import com.openai.auth.X509WorkloadIdentity +import com.openai.core.ClientOptions +import com.openai.core.RequestOptions +import com.openai.core.http.HttpMethod +import com.openai.core.http.HttpRequest +import com.openai.core.http.HttpRequestAttemptAuthenticator +import com.openai.errors.OpenAIIoException +import org.assertj.core.api.Assertions.assertThat +import org.assertj.core.api.Assertions.assertThatThrownBy +import org.junit.jupiter.api.Test + +internal class X509ClientConfigurationTest { + + @Test + fun failedTransportInstallClosesBothLegsAndAuthenticatorButKeepsBuilderReusable() { + val identity = + X509WorkloadIdentity.builder() + .identityProviderId("idp_test") + .serviceAccountId("svc_acct_test") + .build() + val firstExchange = OkHttpClient.builder().build() + val firstApi = OkHttpClient.builder().build() + val firstTransport = BoundX509Transport.create(firstExchange, firstApi) + val failure = IllegalStateException("injected install failure") + var capturedAuthenticator: HttpRequestAttemptAuthenticator? = null + val failing = + X509ClientConfiguration.createForTest(identity, { firstTransport }) { + _, + _, + authenticator -> + capturedAuthenticator = authenticator + throw failure + } + val builder = ClientOptions.builder() + failing.reserve(builder) + + assertThatThrownBy { failing.buildClientOptions(builder) }.isSameAs(failure) + + val request = + HttpRequest.builder().method(HttpMethod.GET).baseUrl(X509_API_BASE_URL).build() + assertThatThrownBy { capturedAuthenticator!!.authenticate(request, null) } + .isInstanceOf(OpenAIIoException::class.java) + .hasMessageContaining("closed") + assertThatThrownBy { firstExchange.execute(request, RequestOptions.none()) } + .isInstanceOf(IllegalStateException::class.java) + .hasMessageContaining("closed") + assertThatThrownBy { firstApi.execute(request, RequestOptions.none()) } + .isInstanceOf(IllegalStateException::class.java) + .hasMessageContaining("closed") + + val secondTransport = + BoundX509Transport.create( + OkHttpClient.builder().build(), + OkHttpClient.builder().build(), + ) + val recovered = + X509ClientConfiguration.create(identity) { secondTransport }.buildClientOptions(builder) + assertThat(recovered.baseUrl()).isEqualTo(X509_API_BASE_URL) + recovered.httpClient.close() + } +} diff --git a/openai-java-core/src/main/kotlin/com/openai/core/ClientOptions.kt b/openai-java-core/src/main/kotlin/com/openai/core/ClientOptions.kt index 169ce7e7b..dc78cdcc5 100644 --- a/openai-java-core/src/main/kotlin/com/openai/core/ClientOptions.kt +++ b/openai-java-core/src/main/kotlin/com/openai/core/ClientOptions.kt @@ -318,26 +318,30 @@ private constructor( explicitBaseUrl = true } - /** Installs the transport owned by a previously reserved fixed-bearer integration. */ + /** + * Builds with a scoped transport owned by a previously reserved fixed-bearer integration. + */ @JvmSynthetic - fun fixedBearerTransport( + fun buildWithFixedBearerTransport( httpClient: HttpClient, httpRequestAuthenticator: HttpRequestAttemptAuthenticator, - ) = apply { - val fixedBaseUrl = - when (val authentication = requestAuthentication) { - is RequestAuthentication.FixedBearerReserved -> - authentication.fixedBearerBaseUrl - is RequestAuthentication.FixedBearerInstalled -> - authentication.fixedBearerBaseUrl - else -> error("Fixed bearer authentication must be set first") - } - this.httpClient = PhantomReachableClosingHttpClient(httpClient) - requestAuthentication = - RequestAuthentication.FixedBearerInstalled.create( - fixedBaseUrl, - httpRequestAuthenticator, - ) + ): ClientOptions { + val originalAuthentication = requestAuthentication + val reserved = originalAuthentication as? RequestAuthentication.FixedBearerReserved + checkNotNull(reserved) { "Fixed bearer authentication must be set first" } + val originalHttpClient = this.httpClient + return try { + this.httpClient = PhantomReachableClosingHttpClient(httpClient) + requestAuthentication = + RequestAuthentication.FixedBearerInstalled.create( + reserved.fixedBearerBaseUrl, + httpRequestAuthenticator, + ) + build() + } finally { + this.httpClient = originalHttpClient + requestAuthentication = originalAuthentication + } } /** diff --git a/openai-java-core/src/main/kotlin/com/openai/core/http/HttpRequestAttemptAuthenticator.kt b/openai-java-core/src/main/kotlin/com/openai/core/http/HttpRequestAttemptAuthenticator.kt index eee12f853..5f4a51c75 100644 --- a/openai-java-core/src/main/kotlin/com/openai/core/http/HttpRequestAttemptAuthenticator.kt +++ b/openai-java-core/src/main/kotlin/com/openai/core/http/HttpRequestAttemptAuthenticator.kt @@ -3,6 +3,33 @@ package com.openai.core.http import java.time.Duration import java.util.concurrent.CompletableFuture +/** Remaining transport timeouts for one authentication attempt. */ +class HttpRequestAttemptTimeouts +private constructor( + private val connect: Duration?, + private val read: Duration?, + private val write: Duration?, + private val request: Duration?, +) { + @JvmSynthetic fun connect(): Duration? = connect + + @JvmSynthetic fun read(): Duration? = read + + @JvmSynthetic fun write(): Duration? = write + + @JvmSynthetic fun request(): Duration? = request + + companion object { + @JvmSynthetic + fun create( + connect: Duration?, + read: Duration?, + write: Duration?, + request: Duration?, + ): HttpRequestAttemptTimeouts = HttpRequestAttemptTimeouts(connect, read, write, request) + } +} + /** * Reserved authentication seam for integrations that must react to the exact request rejected by * the server without starting a second retry lifecycle. @@ -12,6 +39,13 @@ interface HttpRequestAttemptAuthenticator : AutoCloseable { @JvmSynthetic fun authenticate(request: HttpRequest, timeout: Duration?): AuthenticatedHttpRequest + /** Authenticates one attempt using the orchestrator's remaining transport timeouts. */ + @JvmSynthetic + fun authenticate( + request: HttpRequest, + timeouts: HttpRequestAttemptTimeouts, + ): AuthenticatedHttpRequest = authenticate(request, timeouts.request()) + @JvmSynthetic fun authenticateAsync( request: HttpRequest, @@ -25,6 +59,13 @@ interface HttpRequestAttemptAuthenticator : AutoCloseable { } } + /** Authenticates one async attempt using the orchestrator's remaining transport timeouts. */ + @JvmSynthetic + fun authenticateAsync( + request: HttpRequest, + timeouts: HttpRequestAttemptTimeouts, + ): CompletableFuture = authenticateAsync(request, timeouts.request()) + override fun close() {} } diff --git a/openai-java-core/src/main/kotlin/com/openai/core/http/PhantomReachableClosingHttpRequestAttemptAuthenticator.kt b/openai-java-core/src/main/kotlin/com/openai/core/http/PhantomReachableClosingHttpRequestAttemptAuthenticator.kt index fc62de3e3..0af8b8481 100644 --- a/openai-java-core/src/main/kotlin/com/openai/core/http/PhantomReachableClosingHttpRequestAttemptAuthenticator.kt +++ b/openai-java-core/src/main/kotlin/com/openai/core/http/PhantomReachableClosingHttpRequestAttemptAuthenticator.kt @@ -15,11 +15,22 @@ internal class PhantomReachableClosingHttpRequestAttemptAuthenticator( override fun authenticate(request: HttpRequest, timeout: Duration?): AuthenticatedHttpRequest = authenticator.authenticate(request, timeout) + override fun authenticate( + request: HttpRequest, + timeouts: HttpRequestAttemptTimeouts, + ): AuthenticatedHttpRequest = authenticator.authenticate(request, timeouts) + override fun authenticateAsync( request: HttpRequest, timeout: Duration?, ): CompletableFuture = authenticator.authenticateAsync(request, timeout) + override fun authenticateAsync( + request: HttpRequest, + timeouts: HttpRequestAttemptTimeouts, + ): CompletableFuture = + authenticator.authenticateAsync(request, timeouts) + override fun close() = authenticator.close() } diff --git a/openai-java-core/src/main/kotlin/com/openai/core/http/RetryingHttpClientOrchestrator.kt b/openai-java-core/src/main/kotlin/com/openai/core/http/RetryingHttpClientOrchestrator.kt index eaf6f6e2c..d8a6eb064 100644 --- a/openai-java-core/src/main/kotlin/com/openai/core/http/RetryingHttpClientOrchestrator.kt +++ b/openai-java-core/src/main/kotlin/com/openai/core/http/RetryingHttpClientOrchestrator.kt @@ -34,6 +34,7 @@ internal class RetryingHttpClientOrchestrator( private val idempotencyHeader: String?, private val attemptAuthenticator: HttpRequestAttemptAuthenticator?, private val nanoTime: () -> Long = System::nanoTime, + private val beforeAsyncApiDispatch: () -> Unit = {}, ) : HttpClient { private val closed = AtomicBoolean() private val activeAuthenticatedRequests = ConcurrentHashMap.newKeySet>() @@ -127,7 +128,9 @@ internal class RetryingHttpClientOrchestrator( } val authenticated = try { - authenticator.authenticate(current, remainingOrThrow(deadline)) + val authenticationOptions = + deadline?.let { remainingOptions(requestOptions, it) } ?: requestOptions + authenticator.authenticate(current, attemptTimeouts(authenticationOptions)) } catch (error: Throwable) { if (retries >= maxRetries || !shouldRetryAttempt(error)) { throw error @@ -221,7 +224,28 @@ internal class RetryingHttpClientOrchestrator( authenticator: HttpRequestAttemptAuthenticator, ): CompletableFuture { val pipelineRequest = request.withPipelineOwnedBody() - val result = CompletableFuture() + val stageLock = Any() + var terminalReserved = false + val result = + object : CompletableFuture() { + private fun reserveTerminal(): Boolean = + synchronized(stageLock) { + if (terminalReserved || isDone) false + else { + terminalReserved = true + true + } + } + + override fun cancel(mayInterruptIfRunning: Boolean): Boolean = + reserveTerminal() && super.cancel(mayInterruptIfRunning) + + override fun complete(value: HttpResponse): Boolean = + reserveTerminal() && super.complete(value) + + override fun completeExceptionally(error: Throwable): Boolean = + reserveTerminal() && super.completeExceptionally(error) + } val active = AtomicReference?>() result.whenComplete { _, _ -> if (result.isCancelled || closed.get()) active.getAndSet(null)?.cancel(true) @@ -241,10 +265,10 @@ internal class RetryingHttpClientOrchestrator( var retries = 0 var replayed = false - fun activate(future: CompletableFuture<*>) { - active.set(future) - if (result.isDone) future.cancel(true) - } + fun startStage(start: () -> CompletableFuture): CompletableFuture? = + synchronized(stageLock) { + if (terminalReserved || result.isDone) null else start().also { active.set(it) } + } fun closeDiscarded(response: HttpResponse?) { try { @@ -320,12 +344,11 @@ internal class RetryingHttpClientOrchestrator( } val sleep = try { - sleeper.sleepAsync(delay) + startStage { sleeper.sleepAsync(delay) } ?: return } catch (sleepError: Throwable) { result.completeExceptionally(sleepError) return } - activate(sleep) sleep.whenComplete { _, sleepError -> try { if (sleepError == null) run() @@ -343,14 +366,14 @@ internal class RetryingHttpClientOrchestrator( fun dispatch(authenticated: AuthenticatedHttpRequest, options: RequestOptions) { val authenticatedRequest = authenticated.request() + beforeAsyncApiDispatch() val call = try { - httpClient.executeAsync(authenticatedRequest, options) + startStage { httpClient.executeAsync(authenticatedRequest, options) } ?: return } catch (error: Throwable) { retry(error = error, requestRetryable = isRetryable(authenticatedRequest)) return } - activate(call) call.whenComplete callComplete@{ response, callError -> var ownedResponse = response try { @@ -388,9 +411,9 @@ internal class RetryingHttpClientOrchestrator( run = run@{ if (result.isDone) return@run - val timeout = + val authenticationOptions = try { - remainingOrThrow(deadline) + deadline?.let { remainingOptions(requestOptions, it) } ?: requestOptions } catch (error: Throwable) { result.completeExceptionally(error) return@run @@ -398,12 +421,16 @@ internal class RetryingHttpClientOrchestrator( val current = if (sendRetryCount) setRetryCountHeader(modified, retries) else modified val authentication = try { - authenticator.authenticateAsync(current, timeout) + startStage { + authenticator.authenticateAsync( + current, + attemptTimeouts(authenticationOptions), + ) + } ?: return@run } catch (error: Throwable) { retry(error = error, authenticationFailure = true) return@run } - activate(authentication) authentication.whenComplete authenticationComplete@{ authenticated, authError -> try { if (result.isDone) return@authenticationComplete @@ -497,9 +524,6 @@ internal class RetryingHttpClientOrchestrator( return Duration.ofNanos(nanos) } - private fun remainingOrThrow(deadline: Deadline?): Duration? = - deadline?.let(::remaining)?.also { if (it.isZero) throw timedOut() } - private fun remainingOptions(options: RequestOptions, deadline: Deadline): RequestOptions { val remaining = remaining(deadline) if (remaining.isZero) throw timedOut() @@ -508,6 +532,16 @@ internal class RetryingHttpClientOrchestrator( ) } + private fun attemptTimeouts(options: RequestOptions): HttpRequestAttemptTimeouts { + val timeout = options.timeout + return HttpRequestAttemptTimeouts.create( + timeout?.connect(), + timeout?.read(), + timeout?.write(), + timeout?.request(), + ) + } + private fun sleepAuthenticated(delay: Duration, deadline: Deadline?) { val remaining = deadline?.let(::remaining) if (remaining != null && delay >= remaining) throw timedOut() diff --git a/openai-java-core/src/test/kotlin/com/openai/core/CancellationPropagatingFutureTest.kt b/openai-java-core/src/test/kotlin/com/openai/core/CancellationPropagatingFutureTest.kt index 4bb1015f1..7526bed26 100644 --- a/openai-java-core/src/test/kotlin/com/openai/core/CancellationPropagatingFutureTest.kt +++ b/openai-java-core/src/test/kotlin/com/openai/core/CancellationPropagatingFutureTest.kt @@ -377,9 +377,8 @@ internal class CancellationPropagatingFutureTest { OpenAIClientAsyncImpl( ClientOptions.builder() .fixedBearerAuthentication("https://example.test/v1") - .fixedBearerTransport(transport, authenticator) .maxRetries(0) - .build() + .buildWithFixedBearerTransport(transport, authenticator) ) try { @@ -442,9 +441,8 @@ internal class CancellationPropagatingFutureTest { OpenAIClientImpl( ClientOptions.builder() .fixedBearerAuthentication("https://example.test/v1") - .fixedBearerTransport(transport, authenticator) .maxRetries(0) - .build() + .buildWithFixedBearerTransport(transport, authenticator) ) try { @@ -500,9 +498,8 @@ internal class CancellationPropagatingFutureTest { OpenAIClientAsyncImpl( ClientOptions.builder() .fixedBearerAuthentication("https://example.test/v1") - .fixedBearerTransport(transport, authenticator) .maxRetries(0) - .build() + .buildWithFixedBearerTransport(transport, authenticator) ) try { @@ -601,9 +598,8 @@ internal class CancellationPropagatingFutureTest { OpenAIClientAsyncImpl( ClientOptions.builder() .fixedBearerAuthentication("https://example.test/v1") - .fixedBearerTransport(transport, authenticator) .maxRetries(0) - .build() + .buildWithFixedBearerTransport(transport, authenticator) ) try { @@ -656,9 +652,8 @@ internal class CancellationPropagatingFutureTest { OpenAIClientAsyncImpl( ClientOptions.builder() .fixedBearerAuthentication("https://example.test/v1") - .fixedBearerTransport(transport, authenticator) .maxRetries(0) - .build() + .buildWithFixedBearerTransport(transport, authenticator) ) try { @@ -694,9 +689,8 @@ internal class CancellationPropagatingFutureTest { return OpenAIClientAsyncImpl( ClientOptions.builder() .fixedBearerAuthentication("https://example.test/v1") - .fixedBearerTransport(transport, authenticator) .maxRetries(0) - .build() + .buildWithFixedBearerTransport(transport, authenticator) ) } diff --git a/openai-java-core/src/test/kotlin/com/openai/core/ClientOptionsTest.kt b/openai-java-core/src/test/kotlin/com/openai/core/ClientOptionsTest.kt index df0ccfab0..4fe9cceaf 100644 --- a/openai-java-core/src/test/kotlin/com/openai/core/ClientOptionsTest.kt +++ b/openai-java-core/src/test/kotlin/com/openai/core/ClientOptionsTest.kt @@ -193,8 +193,7 @@ internal class ClientOptionsTest { val clientOptions = ClientOptions.builder() .fixedBearerAuthentication("https://mtls.example.test/v1") - .fixedBearerTransport(httpClient, authenticator) - .build() + .buildWithFixedBearerTransport(httpClient, authenticator) .toBuilder() .build() diff --git a/openai-java-core/src/test/kotlin/com/openai/core/X509BlockingResponseLifecycleTest.kt b/openai-java-core/src/test/kotlin/com/openai/core/X509BlockingResponseLifecycleTest.kt index bb3851cc4..2eab95358 100644 --- a/openai-java-core/src/test/kotlin/com/openai/core/X509BlockingResponseLifecycleTest.kt +++ b/openai-java-core/src/test/kotlin/com/openai/core/X509BlockingResponseLifecycleTest.kt @@ -87,9 +87,8 @@ internal class X509BlockingResponseLifecycleTest { } ClientOptions.builder() .fixedBearerAuthentication("https://example.test/v1") - .fixedBearerTransport(transport, authenticator) .maxRetries(0) - .build() + .buildWithFixedBearerTransport(transport, authenticator) } else { ClientOptions.builder() .apiKey("test-api-key") diff --git a/openai-java-core/src/test/kotlin/com/openai/core/http/AttemptAuthenticatingRetryingHttpClientLifecycleTest.kt b/openai-java-core/src/test/kotlin/com/openai/core/http/AttemptAuthenticatingRetryingHttpClientLifecycleTest.kt new file mode 100644 index 000000000..e8d9931a3 --- /dev/null +++ b/openai-java-core/src/test/kotlin/com/openai/core/http/AttemptAuthenticatingRetryingHttpClientLifecycleTest.kt @@ -0,0 +1,120 @@ +package com.openai.core.http + +import com.openai.core.RequestOptions +import com.openai.core.Timeout +import java.time.Clock +import java.time.Duration +import java.util.concurrent.CompletableFuture +import java.util.concurrent.CountDownLatch +import java.util.concurrent.Executors +import java.util.concurrent.TimeUnit +import org.assertj.core.api.Assertions.assertThat +import org.junit.jupiter.api.Test + +internal class AttemptAuthenticatingRetryingHttpClientLifecycleTest { + @Test + fun cancellationWinningAuthenticationToDispatchTransitionPreventsApiCall() { + val authentication = CompletableFuture() + val authenticator = + object : HttpRequestAttemptAuthenticator { + override fun authenticate( + request: HttpRequest, + timeout: Duration?, + ): AuthenticatedHttpRequest = error("sync path not expected") + + override fun authenticateAsync( + request: HttpRequest, + timeout: Duration?, + ): CompletableFuture = authentication + } + val transport = ScriptedClient(response(200)) + val beforeDispatch = CountDownLatch(1) + val allowDispatch = CountDownLatch(1) + val client = + RetryingHttpClientOrchestrator( + httpClient = transport, + sleeper = ImmediateSleeper, + clock = Clock.systemUTC(), + maxRetries = 0, + idempotencyHeader = null, + attemptAuthenticator = authenticator, + beforeAsyncApiDispatch = { + beforeDispatch.countDown() + check(allowDispatch.await(5, TimeUnit.SECONDS)) + }, + ) + val result = client.executeAsync(request()) + val executor = Executors.newSingleThreadExecutor() + + try { + val completion = + executor.submit { + authentication.complete(AuthenticatedHttpRequest.create(request()) {}) + } + assertThat(beforeDispatch.await(5, TimeUnit.SECONDS)).isTrue() + + assertThat(result.cancel(true)).isTrue() + allowDispatch.countDown() + completion.get(5, TimeUnit.SECONDS) + + assertThat(transport.calls).isZero() + } finally { + allowDispatch.countDown() + executor.shutdownNow() + client.close() + } + } + + @Test + fun syncAndAsyncAuthenticationReceiveRemainingRequestOptions() { + listOf(false, true).forEach { async -> + val authenticationTimeouts = mutableListOf() + val authenticator = + object : HttpRequestAttemptAuthenticator { + override fun authenticate( + request: HttpRequest, + timeout: Duration?, + ): AuthenticatedHttpRequest = error("Duration-only boundary must not be used") + + override fun authenticate( + request: HttpRequest, + timeouts: HttpRequestAttemptTimeouts, + ): AuthenticatedHttpRequest { + authenticationTimeouts += timeouts + return AuthenticatedHttpRequest.create(request) {} + } + + override fun authenticateAsync( + request: HttpRequest, + timeouts: HttpRequestAttemptTimeouts, + ): CompletableFuture = + CompletableFuture.completedFuture(authenticate(request, timeouts)) + } + val transport = ScriptedClient(response(200)) + val client = client(transport, authenticator, maxRetries = 0) + val timeout = + Timeout.builder() + .connect(Duration.ofMillis(11)) + .read(Duration.ofMillis(22)) + .write(Duration.ofMillis(33)) + .request(Duration.ofSeconds(5)) + .build() + val options = RequestOptions.builder().timeout(timeout).build() + + if (async) client.executeAsync(request(), options).get(5, TimeUnit.SECONDS).close() + else client.execute(request(), options).close() + + val effective = authenticationTimeouts.single() + assertThat(effective.connect()).isEqualTo(Duration.ofMillis(11)) + assertThat(effective.read()).isEqualTo(Duration.ofMillis(22)) + assertThat(effective.write()).isEqualTo(Duration.ofMillis(33)) + assertThat(effective.request()) + .isNotNull() + .isPositive() + .isLessThanOrEqualTo(Duration.ofSeconds(5)) + assertThat(transport.timeouts.single()) + .isPositive() + .isLessThan(requireNotNull(effective.request())) + } + } +} diff --git a/openai-java-core/src/test/kotlin/com/openai/core/http/PhantomReachableClosingHttpRequestAttemptAuthenticatorTest.kt b/openai-java-core/src/test/kotlin/com/openai/core/http/PhantomReachableClosingHttpRequestAttemptAuthenticatorTest.kt new file mode 100644 index 000000000..3cbc3769f --- /dev/null +++ b/openai-java-core/src/test/kotlin/com/openai/core/http/PhantomReachableClosingHttpRequestAttemptAuthenticatorTest.kt @@ -0,0 +1,49 @@ +package com.openai.core.http + +import java.time.Duration +import java.util.concurrent.CompletableFuture +import java.util.concurrent.TimeUnit +import org.assertj.core.api.Assertions.assertThat +import org.junit.jupiter.api.Test + +internal class PhantomReachableClosingHttpRequestAttemptAuthenticatorTest { + + @Test + fun forwardsAllSyncAndAsyncAttemptTimeouts() { + val observed = mutableListOf() + val delegate = + object : HttpRequestAttemptAuthenticator { + override fun authenticate( + request: HttpRequest, + timeout: Duration?, + ): AuthenticatedHttpRequest = error("Duration-only boundary must not be used") + + override fun authenticate( + request: HttpRequest, + timeouts: HttpRequestAttemptTimeouts, + ): AuthenticatedHttpRequest { + observed += timeouts + return AuthenticatedHttpRequest.create(request) {} + } + + override fun authenticateAsync( + request: HttpRequest, + timeouts: HttpRequestAttemptTimeouts, + ): CompletableFuture = + CompletableFuture.completedFuture(authenticate(request, timeouts)) + } + val wrapper = PhantomReachableClosingHttpRequestAttemptAuthenticator(delegate) + val timeouts = + HttpRequestAttemptTimeouts.create( + Duration.ofMillis(11), + Duration.ofMillis(22), + Duration.ofMillis(33), + Duration.ofMillis(44), + ) + + wrapper.authenticate(request(), timeouts) + wrapper.authenticateAsync(request(), timeouts).get(5, TimeUnit.SECONDS) + + assertThat(observed).containsExactly(timeouts, timeouts) + } +} diff --git a/openai-java-core/src/test/kotlin/com/openai/core/http/X509AsyncStreamCancellationTest.kt b/openai-java-core/src/test/kotlin/com/openai/core/http/X509AsyncStreamCancellationTest.kt index 335903ec1..cb0724d7e 100644 --- a/openai-java-core/src/test/kotlin/com/openai/core/http/X509AsyncStreamCancellationTest.kt +++ b/openai-java-core/src/test/kotlin/com/openai/core/http/X509AsyncStreamCancellationTest.kt @@ -133,9 +133,8 @@ internal class X509AsyncStreamCancellationTest { OpenAIClientAsyncImpl( ClientOptions.builder() .fixedBearerAuthentication("https://example.test/v1") - .fixedBearerTransport(transport, authenticator) .maxRetries(0) - .build() + .buildWithFixedBearerTransport(transport, authenticator) ) private fun params() = From 3bd700aaa195bf2337218645d8b9d738be984ce8 Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Wed, 26 Aug 2026 00:06:39 +0000 Subject: [PATCH 4/5] fix(client): close request bodies on rejected calls --- .../com/openai/client/okhttp/OkHttpClient.kt | 39 +++++++++----- .../openai/client/okhttp/OkHttpClientTest.kt | 54 +++++++++++++++++++ 2 files changed, 81 insertions(+), 12 deletions(-) diff --git a/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/OkHttpClient.kt b/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/OkHttpClient.kt index 160705a41..dd4c845d5 100644 --- a/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/OkHttpClient.kt +++ b/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/OkHttpClient.kt @@ -97,19 +97,24 @@ private constructor( } private fun newCall(request: HttpRequest, requestOptions: RequestOptions): Call { - callTracker.ensureOpen() - val clientBuilder = okHttpClient.newBuilder() - - requestOptions.timeout?.let { - clientBuilder - .connectTimeout(it.connect()) - .readTimeout(it.read()) - .writeTimeout(it.write()) - .callTimeout(it.request()) - } + return try { + callTracker.ensureOpen() + val clientBuilder = okHttpClient.newBuilder() + + requestOptions.timeout?.let { + clientBuilder + .connectTimeout(it.connect()) + .readTimeout(it.read()) + .writeTimeout(it.write()) + .callTimeout(it.request()) + } - val client = clientBuilder.build() - return client.newCall(request.toRequest(client)) + val client = clientBuilder.build() + client.newCall(request.toRequest(client)) + } catch (failure: Throwable) { + request.body.closeSuppressing(failure) + throw failure + } } companion object { @@ -252,6 +257,16 @@ private constructor( } } +private fun HttpRequestBody?.closeSuppressing(failure: Throwable) { + try { + this?.close() + } catch (closeFailure: Throwable) { + if (closeFailure !== failure) { + failure.addSuppressed(closeFailure) + } + } +} + private class CallTracker { private val closed = AtomicBoolean() private val activeCalls = ConcurrentHashMap.newKeySet() diff --git a/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/OkHttpClientTest.kt b/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/OkHttpClientTest.kt index f22297b9e..dabd8b093 100644 --- a/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/OkHttpClientTest.kt +++ b/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/OkHttpClientTest.kt @@ -6,8 +6,10 @@ import com.github.tomakehurst.wiremock.junit5.WireMockTest import com.openai.core.http.Headers import com.openai.core.http.HttpMethod import com.openai.core.http.HttpRequest +import com.openai.core.http.HttpRequestBody import com.openai.core.http.HttpResponse import java.io.ByteArrayInputStream +import java.io.OutputStream import java.util.concurrent.CompletableFuture import java.util.concurrent.CountDownLatch import java.util.concurrent.ExecutionException @@ -17,6 +19,7 @@ import okhttp3.mockwebserver.MockResponse import okhttp3.mockwebserver.MockWebServer import okhttp3.mockwebserver.SocketPolicy import org.assertj.core.api.Assertions.assertThat +import org.assertj.core.api.Assertions.assertThatThrownBy import org.junit.jupiter.api.BeforeEach import org.junit.jupiter.api.Test import org.junit.jupiter.api.parallel.ResourceLock @@ -53,6 +56,32 @@ internal class OkHttpClientTest { assertThat(call.isCanceled()).isTrue() } + @Test + fun execute_afterClientClose_closesRequestBodyOnce() { + val closeFailure = IllegalStateException("request body close failed") + val body = CountingRequestBody(closeFailure) + httpClient.close() + + val failure = runCatching { httpClient.execute(request(body)) }.exceptionOrNull() + + assertThat(failure) + .isInstanceOf(IllegalStateException::class.java) + .hasMessage("HTTP client is closed") + assertThat(failure!!.suppressed).containsExactly(closeFailure) + assertThat(body.closes).isEqualTo(1) + } + + @Test + fun executeAsync_afterClientClose_closesRequestBodyOnce() { + val body = CountingRequestBody() + httpClient.close() + + assertThatThrownBy { httpClient.executeAsync(request(body)) } + .isInstanceOf(IllegalStateException::class.java) + .hasMessage("HTTP client is closed") + assertThat(body.closes).isEqualTo(1) + } + @Test fun completeOrCloseResponse_whenCancellationWins_closesTheDroppedResponse() { val future = CompletableFuture() @@ -123,6 +152,31 @@ internal class OkHttpClientTest { } } } + + private fun request(body: HttpRequestBody): HttpRequest = + HttpRequest.builder() + .method(HttpMethod.POST) + .baseUrl(baseUrl) + .addPathSegment("something") + .body(body) + .build() +} + +private class CountingRequestBody(private val closeFailure: Throwable? = null) : HttpRequestBody { + var closes = 0 + + override fun writeTo(outputStream: OutputStream) {} + + override fun contentType(): String? = null + + override fun contentLength(): Long = 0 + + override fun repeatable(): Boolean = true + + override fun close() { + closes++ + closeFailure?.let { throw it } + } } private class TrackingHttpResponse : HttpResponse { From 9a3ac279fa8a56edc9df884f923c777707c105c7 Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Wed, 26 Aug 2026 00:14:08 +0000 Subject: [PATCH 5/5] fix(client): preserve transport failures during cleanup --- .../com/openai/client/okhttp/OkHttpClient.kt | 15 +++++++++-- .../openai/client/okhttp/OkHttpClientTest.kt | 26 +++++++++++++++++-- 2 files changed, 37 insertions(+), 4 deletions(-) diff --git a/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/OkHttpClient.kt b/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/OkHttpClient.kt index dd4c845d5..8f8f4850c 100644 --- a/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/OkHttpClient.kt +++ b/openai-java-client-okhttp/src/main/kotlin/com/openai/client/okhttp/OkHttpClient.kt @@ -50,12 +50,23 @@ private constructor( override fun execute(request: HttpRequest, requestOptions: RequestOptions): HttpResponse { val call = newCall(request, requestOptions) + var failure: Throwable? = null return try { call.execute().toHttpResponse() } catch (e: IOException) { - throw OpenAIIoException("Request failed", e) + val ioFailure = OpenAIIoException("Request failed", e) + failure = ioFailure + throw ioFailure + } catch (t: Throwable) { + failure = t + throw t } finally { - request.body?.close() + val primaryFailure = failure + if (primaryFailure == null) { + request.body?.close() + } else { + request.body.closeSuppressing(primaryFailure) + } } } diff --git a/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/OkHttpClientTest.kt b/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/OkHttpClientTest.kt index dabd8b093..585064186 100644 --- a/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/OkHttpClientTest.kt +++ b/openai-java-client-okhttp/src/test/kotlin/com/openai/client/okhttp/OkHttpClientTest.kt @@ -8,6 +8,7 @@ import com.openai.core.http.HttpMethod import com.openai.core.http.HttpRequest import com.openai.core.http.HttpRequestBody import com.openai.core.http.HttpResponse +import com.openai.errors.OpenAIIoException import java.io.ByteArrayInputStream import java.io.OutputStream import java.util.concurrent.CompletableFuture @@ -82,6 +83,27 @@ internal class OkHttpClientTest { assertThat(body.closes).isEqualTo(1) } + @Test + fun execute_transportFailureSuppressesRequestBodyCloseFailure() { + val server = MockWebServer() + val closeFailure = IllegalStateException("request body close failed") + val body = CountingRequestBody(closeFailure) + try { + server.start() + server.enqueue(MockResponse().setSocketPolicy(SocketPolicy.DISCONNECT_AT_START)) + + val failure = + runCatching { httpClient.execute(request(body, server.url("/").toString())) } + .exceptionOrNull() + + assertThat(failure).isInstanceOf(OpenAIIoException::class.java) + assertThat(failure!!.suppressed).containsExactly(closeFailure) + assertThat(body.closes).isEqualTo(1) + } finally { + server.close() + } + } + @Test fun completeOrCloseResponse_whenCancellationWins_closesTheDroppedResponse() { val future = CompletableFuture() @@ -153,10 +175,10 @@ internal class OkHttpClientTest { } } - private fun request(body: HttpRequestBody): HttpRequest = + private fun request(body: HttpRequestBody, requestBaseUrl: String = baseUrl): HttpRequest = HttpRequest.builder() .method(HttpMethod.POST) - .baseUrl(baseUrl) + .baseUrl(requestBaseUrl) .addPathSegment("something") .body(body) .build()