diff --git a/anthropic-java-core/src/main/kotlin/com/anthropic/core/http/AsyncStreamResponse.kt b/anthropic-java-core/src/main/kotlin/com/anthropic/core/http/AsyncStreamResponse.kt index 5bb2e37c4..a97ee965b 100644 --- a/anthropic-java-core/src/main/kotlin/com/anthropic/core/http/AsyncStreamResponse.kt +++ b/anthropic-java-core/src/main/kotlin/com/anthropic/core/http/AsyncStreamResponse.kt @@ -94,46 +94,64 @@ internal fun CompletableFuture>.toAsync(streamHandlerExecu else "Cannot subscribe after the response is closed" } - this@toAsync.whenCompleteAsync( - { streamResponse, futureError -> - if (state.get() == State.CLOSED) { - // Avoid doing any work if `close` was called before the future - // completed. - return@whenCompleteAsync - } + val completion = + this@toAsync.whenCompleteAsync( + { streamResponse, futureError -> + if (state.get() == State.CLOSED) { + // Avoid doing any work if `close` was called before the future + // completed. + return@whenCompleteAsync + } - if (futureError != null) { - // An error occurred before we started passing chunks to the handler. - handler.onComplete(Optional.of(futureError)) - return@whenCompleteAsync - } + if (futureError != null) { + // An error occurred before we started passing chunks to the + // handler. + handler.onComplete(Optional.of(futureError)) + return@whenCompleteAsync + } - var streamError: Throwable? = null - try { - streamResponse.stream().forEach(handler::onNext) - } catch (e: Throwable) { - streamError = e - } + var streamError: Throwable? = null + try { + streamResponse.stream().forEach(handler::onNext) + } catch (e: Throwable) { + streamError = e + } - try { - handler.onComplete(Optional.ofNullable(streamError)) - } finally { try { - // Notify completion via the `onCompleteFuture` as well. This is in - // a separate `try-finally` block so that we still complete the - // future if `handler.onComplete` throws. - if (streamError == null) { - onCompleteFuture.complete(null) - } else { - onCompleteFuture.completeExceptionally(streamError) - } + handler.onComplete(Optional.ofNullable(streamError)) } finally { - close() + try { + // Notify completion via the `onCompleteFuture` as well. This is + // in + // a separate `try-finally` block so that we still complete the + // future if `handler.onComplete` throws. + if (streamError == null) { + onCompleteFuture.complete(null) + } else { + onCompleteFuture.completeExceptionally(streamError) + } + } finally { + close() + } } + }, + executor, + ) + + // Bookkeeping must still run if the callback executor rejects work. + // As with AutoPagerAsync, never move user callbacks to another thread. + @Suppress("ForbiddenMethodCall") + completion.whenComplete { _, error -> + if (state.get() != State.CLOSED && error != null) { + // The response future is already complete here. Preserve its + // original failure even if rejection raced its bookkeeping. + @Suppress("ForbiddenMethodCall") + this@toAsync.whenComplete { _, responseError -> + onCompleteFuture.completeExceptionally(responseError ?: error) + close() } - }, - executor, - ) + } + } } override fun onCompleteFuture(): CompletableFuture = onCompleteFuture diff --git a/anthropic-java-core/src/test/kotlin/com/anthropic/core/http/AsyncStreamResponseTest.kt b/anthropic-java-core/src/test/kotlin/com/anthropic/core/http/AsyncStreamResponseTest.kt index b24af2310..0cc2b2b60 100644 --- a/anthropic-java-core/src/test/kotlin/com/anthropic/core/http/AsyncStreamResponseTest.kt +++ b/anthropic-java-core/src/test/kotlin/com/anthropic/core/http/AsyncStreamResponseTest.kt @@ -2,7 +2,11 @@ package com.anthropic.core.http import java.util.* import java.util.concurrent.CompletableFuture +import java.util.concurrent.CompletionException import java.util.concurrent.Executor +import java.util.concurrent.Executors +import java.util.concurrent.RejectedExecutionException +import java.util.concurrent.TimeUnit import java.util.stream.Stream import kotlin.streams.asStream import org.assertj.core.api.Assertions.assertThat @@ -265,4 +269,129 @@ internal class AsyncStreamResponseTest { assertDoesNotThrow { future.completeExceptionally(ERROR) } } + + @Test + fun rejectedDefaultExecutor_finishesExceptionallyAndClosesUnreadResponse() { + for (alreadyCompleted in listOf(false, true)) { + val pool = Executors.newSingleThreadExecutor() + pool.shutdown() + val response = TrackingResponse() + val responseFuture = CompletableFuture>() + if (alreadyCompleted) responseFuture.complete(response) + val asyncResponse = responseFuture.toAsync(pool) + val callback = mock>() + try { + assertDoesNotThrow { asyncResponse.subscribe(callback) } + responseFuture.complete(response) + assertThat(asyncResponse.onCompleteFuture()).isCompletedExceptionally + val failure = catchThrowable { asyncResponse.onCompleteFuture().join() } + assertThat(failure) + .isInstanceOf(CompletionException::class.java) + .hasCauseInstanceOf(RejectedExecutionException::class.java) + assertThat(response.closes).isEqualTo(1) + assertThat(response.streamCalls).isZero() + verifyNoInteractions(callback) + asyncResponse.close() + assertThat(response.closes).isEqualTo(1) + } finally { + asyncResponse.close() + pool.shutdownNow() + } + } + } + + @Test + fun rejectedSubscriptionExecutor_overridesAcceptingDefaultAndReleasesResponse() { + val rejected = RejectedExecutionException("handler executor shut down") + val response = TrackingResponse() + val responseFuture = CompletableFuture>() + val asyncResponse = responseFuture.toAsync(executor) + val callback = mock>() + try { + asyncResponse.subscribe(callback, Executor { throw rejected }) + responseFuture.complete(response) + assertThat(asyncResponse.onCompleteFuture()).isCompletedExceptionally + val failure = catchThrowable { asyncResponse.onCompleteFuture().join() } + assertThat(failure).hasCause(rejected) + assertThat(response.closes).isEqualTo(1) + assertThat(response.streamCalls).isZero() + verifyNoInteractions(callback) + verify(executor, never()).execute(any()) + } finally { + asyncResponse.close() + } + } + + @Test + fun rejectedExecutor_doesNotReplaceAnEarlierTransportFailureOrExplicitClose() { + val rejected = Executor { throw RejectedExecutionException("not accepting work") } + val failedFuture = CompletableFuture>() + val failed = failedFuture.toAsync(rejected) + failed.subscribe(handler) + failedFuture.completeExceptionally(ERROR) + assertThat(catchThrowable { failed.onCompleteFuture().join() }).hasCause(ERROR) + failed.close() + + val response = TrackingResponse() + val pending = CompletableFuture>() + val closed = pending.toAsync(rejected) + closed.subscribe(handler) + closed.close() + pending.complete(response) + assertThat(closed.onCompleteFuture()).isCompleted + assertThat(closed.onCompleteFuture()).isNotCompletedExceptionally + assertThat(response.closes).isEqualTo(1) + assertThat(response.streamCalls).isZero() + verifyNoInteractions(handler) + } + + @Test + fun acceptedExecutor_stillDeliversAndClosesOnItsOwnThread() { + val pool = + Executors.newSingleThreadExecutor { runnable -> Thread(runnable, "stream-worker") } + val response = TrackingResponse() + val pending = CompletableFuture>() + val asyncResponse = pending.toAsync(pool) + val threads = mutableListOf() + val values = mutableListOf() + try { + asyncResponse.subscribe( + object : AsyncStreamResponse.Handler { + override fun onNext(value: String) { + threads.add(Thread.currentThread().name) + values.add(value) + } + + override fun onComplete(error: Optional) { + assertThat(error).isEmpty + threads.add(Thread.currentThread().name) + } + } + ) + pending.complete(response) + asyncResponse.onCompleteFuture().get(5, TimeUnit.SECONDS) + pool.shutdown() + assertThat(pool.awaitTermination(5, TimeUnit.SECONDS)).isTrue() + assertThat(values).containsExactly("one", "two") + assertThat(threads).containsExactly("stream-worker", "stream-worker", "stream-worker") + assertThat(response.closes).isEqualTo(1) + } finally { + asyncResponse.close() + pool.shutdownNow() + } + } + + private class TrackingResponse : StreamResponse { + var closes = 0 + var streamCalls = 0 + + override fun stream(): Stream { + streamCalls++ + return Stream.of("one", "two") + } + + override fun close() { + closes++ + } + } }