Skip to content

fix: clear WorkloadIdentityAuth refresh state after synchronous failures #854

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
abhinavkr26104 wants to merge 3 commits into
base: main
Choose a base branch
from
Open

fix: clear WorkloadIdentityAuth refresh state after synchronous failures #854

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ffb0468
docs: fix async pagination examples
abhinavkr26104 Aug 10, 2026
183fd87
fix: preserve shared resources across derived clients
abhinavkr26104 Aug 10, 2026
a803854
fix: clear workload identity refresh state on sync failures
abhinavkr26104 Aug 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1343,12 +1343,12 @@ import java.util.concurrent.CompletableFuture;

CompletableFuture<JobListPageAsync> pageFuture = client.async().fineTuning().jobs().list();

pageFuture.thenRun(page -> page.autoPager().subscribe(job -> {
pageFuture.thenAccept(page -> page.autoPager().subscribe(job -> {
System.out.println(job);
}));

// If you need to handle errors or completion of the stream
pageFuture.thenRun(page -> page.autoPager().subscribe(new AsyncStreamResponse.Handler<>() {
pageFuture.thenAccept(page -> page.autoPager().subscribe(new AsyncStreamResponse.Handler<>() {
@Override
public void onNext(FineTuningJob job) {
System.out.println(job);
Expand All @@ -1366,7 +1366,7 @@ pageFuture.thenRun(page -> page.autoPager().subscribe(new AsyncStreamResponse.Ha
}));

// Or use futures
pageFuture.thenRun(page -> page.autoPager()
pageFuture.thenAccept(page -> page.autoPager()
.subscribe(job -> {
System.out.println(job);
})
Expand Down
12 changes: 8 additions & 4 deletions openai-java-core/src/main/kotlin/com/openai/auth/WorkloadIdentityAuth.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -236,11 +236,15 @@ internal class WorkloadIdentityAuth(
}

private fun refreshTokenAsync(): CompletableFuture<String> {
return config.provider.getTokenAsync(httpClient, jsonMapper).thenCompose { subjectToken ->
val request = buildTokenExchangeRequest(subjectToken)
httpClient.executeAsync(request).thenApply { response ->
response.use { processTokenExchangeResponse(it) }
return try {
config.provider.getTokenAsync(httpClient, jsonMapper).thenCompose { subjectToken ->
val request = buildTokenExchangeRequest(subjectToken)
httpClient.executeAsync(request).thenApply { response ->
response.use { processTokenExchangeResponse(it) }
}
}
} catch (e: Exception) {
CompletableFuture<String>().also { it.completeExceptionally(e) }
}
}

Expand Down
5 changes: 4 additions & 1 deletion openai-java-core/src/main/kotlin/com/openai/client/OpenAIClientAsyncImpl.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -234,7 +234,10 @@ class OpenAIClientAsyncImpl(private val clientOptions: ClientOptions) : OpenAICl

override fun videos(): VideoServiceAsync = videos

override fun close() = clientOptions.close()
override fun close() {
clientOptionsWithUserAgent.close()
if (clientOptionsWithUserAgent !== clientOptions) clientOptions.close()
}

class WithRawResponseImpl internal constructor(private val clientOptions: ClientOptions) :
OpenAIClientAsync.WithRawResponse {
Expand Down
5 changes: 4 additions & 1 deletion openai-java-core/src/main/kotlin/com/openai/client/OpenAIClientImpl.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -213,7 +213,10 @@ class OpenAIClientImpl(private val clientOptions: ClientOptions) : OpenAIClient

override fun videos(): VideoService = videos

override fun close() = clientOptions.close()
override fun close() {
clientOptionsWithUserAgent.close()
if (clientOptionsWithUserAgent !== clientOptions) clientOptions.close()
}

class WithRawResponseImpl internal constructor(private val clientOptions: ClientOptions) :
OpenAIClient.WithRawResponse {
Expand Down
107 changes: 101 additions & 6 deletions openai-java-core/src/main/kotlin/com/openai/core/ClientOptions.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,9 +30,59 @@ import java.util.concurrent.Executor
import java.util.concurrent.ExecutorService
import java.util.concurrent.Executors
import java.util.concurrent.ThreadFactory
import java.util.concurrent.atomic.AtomicBoolean
import java.util.concurrent.atomic.AtomicLong
import kotlin.jvm.optionals.getOrNull

private class ClientOptionsResource(private val close: () -> Unit) {
private var references = 1
private var closed = false

@Synchronized
fun retain() {
check(!closed) { "Cannot retain a closed client resource" }
references++
}

fun release() {
val shouldClose =
synchronized(this) {
check(references > 0) { "Client resource released too many times" }
references--
if (references == 0) {
closed = true
true
} else false
}

if (shouldClose) close()
}
}

private class ClientOptionsResources(
val httpClient: ClientOptionsResource,
val httpRequestAuthenticator: ClientOptionsResource?,
val workloadIdentityAuth: ClientOptionsResource?,
val streamHandlerExecutor: ClientOptionsResource,
val sleeper: ClientOptionsResource,
) {
fun release() {
httpRequestAuthenticator?.release()
workloadIdentityAuth?.release()
httpClient.release()
streamHandlerExecutor.release()
sleeper.release()
}
}

private class ClientOptionsCloseAction(private val resources: ClientOptionsResources) : () -> Unit {
private val closed = AtomicBoolean(false)

override fun invoke() {
if (closed.compareAndSet(false, true)) resources.release()
}
}

/** A class representing the SDK client configuration. */
class ClientOptions
private constructor(
Expand Down Expand Up @@ -141,12 +191,16 @@ private constructor(
private val organization: String?,
private val project: String?,
private val webhookSecret: String?,
private val resources: ClientOptionsResources,
) {

private val closeAction = ClientOptionsCloseAction(resources)

init {
if (checkJacksonVersionCompatibility) {
checkJacksonVersionCompatibility()
}
closeWhenPhantomReachable(this, closeAction)

@chatgpt-codex-connector chatgpt-codex-connector Bot Aug 10, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Keep resources alive for outstanding async work

If a caller starts an async request or obtains an AsyncStreamResponse and then drops the client/service before the future or stream is finished, this cleaner can run solely because the ClientOptions object is unreachable. The returned async work still uses the underlying HTTP client, retry sleeper, or stream-handler executor, but those are now closed by closeAction; for example AsyncStreamResponse.toAsync() submits stream consumption to clientOptions.streamHandlerExecutor, so shutting it down before subscribe() can leave the stream callbacks unscheduled. Tie cleanup to the outstanding returned work as well, or avoid closing shared resources from the ClientOptions cleaner while such objects may still own them.

Useful? React with 👍 / 👎.

}

/**
Expand Down Expand Up @@ -217,6 +271,11 @@ private constructor(
private var project: String? = null
private var webhookSecret: String? = null
private var workloadIdentity: WorkloadIdentity? = null
private var httpClientResource: ClientOptionsResource? = null
private var httpRequestAuthenticatorResource: ClientOptionsResource? = null
private var workloadIdentityAuthResource: ClientOptionsResource? = null
private var streamHandlerExecutorResource: ClientOptionsResource? = null
private var sleeperResource: ClientOptionsResource? = null

@JvmSynthetic
internal fun from(clientOptions: ClientOptions) = apply {
Expand All @@ -226,6 +285,11 @@ private constructor(
jsonMapper = clientOptions.jsonMapper
streamHandlerExecutor = clientOptions.streamHandlerExecutor
sleeper = clientOptions.sleeper
httpClientResource = clientOptions.resources.httpClient
httpRequestAuthenticatorResource = clientOptions.resources.httpRequestAuthenticator
workloadIdentityAuthResource = clientOptions.resources.workloadIdentityAuth
streamHandlerExecutorResource = clientOptions.resources.streamHandlerExecutor
sleeperResource = clientOptions.resources.sleeper
Comment on lines +288 to +292

@chatgpt-codex-connector chatgpt-codex-connector Bot Aug 10, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Retain resources while cloned builders are alive

When a builder returned by toBuilder() is kept after the source ClientOptions becomes unreachable, these copied resource handles are not retained until build() runs. The new cleaner on ClientOptions can therefore release and close the shared handles while the builder is still alive; a later builder.build() will either throw Cannot retain a closed client resource or clone options backed by resources that have already been closed. Retain these handles in from() and release them if the builder replaces them, or otherwise make the builder keep the source options/resources alive until build time.

Useful? React with 👍 / 👎.

clock = clientOptions.clock
baseUrl = clientOptions.baseUrl
headers = clientOptions.headers.toBuilder()
Expand Down Expand Up @@ -256,6 +320,7 @@ private constructor(
*/
fun httpClient(httpClient: HttpClient) = apply {
this.httpClient = PhantomReachableClosingHttpClient(httpClient)
this.httpClientResource = null
}

/**
Expand All @@ -269,6 +334,7 @@ private constructor(
this.httpRequestAuthenticator =
if (httpRequestAuthenticator == null) null
else PhantomReachableClosingHttpRequestAuthenticator(httpRequestAuthenticator)
this.httpRequestAuthenticatorResource = null
}

/**
Expand Down Expand Up @@ -302,6 +368,7 @@ private constructor(
if (streamHandlerExecutor is ExecutorService)
PhantomReachableExecutorService(streamHandlerExecutor)
else streamHandlerExecutor
this.streamHandlerExecutorResource = null
}

/**
Expand All @@ -313,7 +380,10 @@ private constructor(
*
* This class takes ownership of the sleeper and closes it when closed.
*/
fun sleeper(sleeper: Sleeper) = apply { this.sleeper = PhantomReachableSleeper(sleeper) }
fun sleeper(sleeper: Sleeper) = apply {
this.sleeper = PhantomReachableSleeper(sleeper)
this.sleeperResource = null
}

/**
* The clock to use for operations that require timing, like retries.
Expand Down Expand Up @@ -393,6 +463,7 @@ private constructor(
fun apiKey(apiKey: String?) = apply {
this.apiKey = apiKey
this.credential = apiKey?.let { BearerTokenCredential.create(it) }
this.workloadIdentityAuthResource = null
}

/** Alias for calling [Builder.apiKey] with `apiKey.orElse(null)`. */
Expand All @@ -406,6 +477,7 @@ private constructor(
fun credential(credential: Credential) = apply {
this.apiKey = null
this.credential = credential
this.workloadIdentityAuthResource = null
}

fun azureServiceVersion(azureServiceVersion: AzureOpenAIServiceVersion) = apply {
Expand Down Expand Up @@ -434,6 +506,7 @@ private constructor(

fun workloadIdentity(workloadIdentity: WorkloadIdentity?) = apply {
this.workloadIdentity = workloadIdentity
this.workloadIdentityAuthResource = null

@chatgpt-codex-connector chatgpt-codex-connector Bot Aug 10, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Retain workload identity auth when nulling config

When this builder was created from options that use workload identity, from() keeps the existing WorkloadIdentityCredential in credential. If a caller then does toBuilder().workloadIdentity(null).build(), this line drops the retained workloadIdentityAuthResource even though the old credential is still selected, so the derived options create a fresh owner for the same WorkloadIdentityAuth; closing that derived options/client can close the shared provider, such as the K8s token reader executor, while the original options/client is still active. Clear the workload-identity credential when nulling the config, or keep retaining the existing resource when the credential is unchanged.

Useful? React with 👍 / 👎.

}

/** Alias for calling [Builder.workloadIdentity] with `workloadIdentity.orElse(null)`. */
Expand Down Expand Up @@ -701,6 +774,31 @@ private constructor(
val effectiveWorkloadIdentityAuth =
(credential as? WorkloadIdentityCredential)?.getAuth()

val resources =
ClientOptionsResources(
httpClient =
httpClientResource?.also { it.retain() }
?: ClientOptionsResource { httpClient.close() },
httpRequestAuthenticator =
httpRequestAuthenticatorResource?.also { it.retain() }
?: httpRequestAuthenticator?.let {
ClientOptionsResource { it.close() }
},
workloadIdentityAuth =
workloadIdentityAuthResource?.also { it.retain() }
?: effectiveWorkloadIdentityAuth?.let {
ClientOptionsResource { it.close() }
},
streamHandlerExecutor =
streamHandlerExecutorResource?.also { it.retain() }
?: ClientOptionsResource {
(streamHandlerExecutor as? ExecutorService)?.shutdown()
},
sleeper =
sleeperResource?.also { it.retain() }
?: ClientOptionsResource { sleeper.close() },
)

val loggingDelegate =
if (httpRequestAuthenticator != null) httpClient
else
Expand Down Expand Up @@ -756,6 +854,7 @@ private constructor(
organization,
project,
webhookSecret,
resources,
)
}
}
Expand All @@ -770,11 +869,7 @@ private constructor(
* releases threads and connections if they remain idle, but if you are writing an application
* that needs to aggressively release unused resources, then you may call this method.
*/
fun close() {
httpClient.close()
(streamHandlerExecutor as? ExecutorService)?.shutdown()
sleeper.close()
}
fun close() = closeAction()

@JvmSynthetic
internal fun securityHeaders(security: SecurityOptions): Headers {
Expand Down
Loading