diff --git a/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java b/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java index 43faef5cf..903242ede 100644 --- a/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java +++ b/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java @@ -55,6 +55,7 @@ import io.reactivex.rxjava3.core.Single; import io.reactivex.rxjava3.disposables.Disposable; import io.reactivex.rxjava3.observers.DisposableCompletableObserver; +import io.reactivex.rxjava3.subjects.CompletableSubject; import java.util.ArrayList; import java.util.List; import java.util.Optional; @@ -542,7 +543,8 @@ private Flowable runStepAndCheckContinue( InvocationContext invocationContext, int stepsCompleted, AtomicBoolean continueFlow) { - Flowable currentStepEvents = runOneStep(spanContext, invocationContext).cache(); + Flowable currentStepEvents = + cancellableCache(runOneStep(spanContext, invocationContext)); if (stepsCompleted + 1 >= maxSteps) { logger.debug("Ending flow execution because max steps reached."); return currentStepEvents; @@ -576,6 +578,16 @@ private Flowable runStepAndCheckContinue( })); } + /** + * Like {@code step.cache()}, but cancels {@code step} when a subscriber cancels. {@code cache()} + * never cancels upstream, so a client disconnect left the model stream and pending tool calls + * running. A subscriber that arrives after the cancel still gets the buffered events. + */ + private static Flowable cancellableCache(Flowable step) { + CompletableSubject cancelled = CompletableSubject.create(); + return step.takeUntil(cancelled.toFlowable()).cache().doOnCancel(cancelled::onComplete); + } + /** * Executes the LLM flow in streaming mode. * diff --git a/core/src/main/java/com/google/adk/models/Gemini.java b/core/src/main/java/com/google/adk/models/Gemini.java index 09402c2b5..afb50be9c 100644 --- a/core/src/main/java/com/google/adk/models/Gemini.java +++ b/core/src/main/java/com/google/adk/models/Gemini.java @@ -38,14 +38,17 @@ import com.google.genai.types.Part; import com.google.genai.types.PartialArg; import io.reactivex.rxjava3.core.Flowable; +import io.reactivex.rxjava3.schedulers.Schedulers; import java.time.Duration; import java.util.ArrayList; +import java.util.Iterator; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Objects; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutorService; +import java.util.concurrent.atomic.AtomicBoolean; import okhttp3.OkHttpClient; import org.jspecify.annotations.Nullable; import org.slf4j.Logger; @@ -299,11 +302,7 @@ private Flowable sendRequest( effectiveModelName, llmRequest.contents(), config); return Flowable.defer( - () -> - new StreamingResponseAggregator() - .process( - Flowable.fromFuture(streamFuture).flatMapIterable(iterable -> iterable), - cacheMetadata)); + () -> new StreamingResponseAggregator().process(readStream(streamFuture), cacheMetadata)); } else { logger.debug("Sending generateContent request to model {}", effectiveModelName); return Flowable.fromFuture( @@ -316,6 +315,78 @@ private Flowable sendRequest( } } + /** + * Reads the stream and, when the subscriber cancels, closes it on {@link Schedulers#io()} so a + * partial read doesn't leave the HTTP response open. The close waits for the stream to open and + * can block behind a {@code readLine()} in progress. An iterator call that returns after the + * cancel closes the stream again, covering a continuation request opened during that call. + */ + @VisibleForTesting + static & AutoCloseable> Flowable readStream( + CompletableFuture streamFuture) { + return Flowable.defer( + () -> { + AtomicBoolean cancelled = new AtomicBoolean(); + AtomicBoolean terminated = new AtomicBoolean(); + return Flowable.using( + () -> streamFuture, + future -> + Flowable.fromFuture(future) + .flatMapIterable(stream -> closingAfterCancel(stream, cancelled)) + .doOnTerminate(() -> terminated.set(true)), + future -> { + if (!terminated.get()) { + cancelled.set(true); + future.thenAccept(Gemini::closeAsync); + } + }, + /* eager= */ false); + }); + } + + private static & AutoCloseable> Iterable closingAfterCancel( + S stream, AtomicBoolean cancelled) { + Iterator delegate = stream.iterator(); + return () -> + new Iterator() { + @Override + public boolean hasNext() { + try { + return delegate.hasNext(); + } finally { + closeIfCancelled(); + } + } + + @Override + public T next() { + try { + return delegate.next(); + } finally { + closeIfCancelled(); + } + } + + private void closeIfCancelled() { + if (cancelled.get()) { + closeAsync(stream); + } + } + }; + } + + private static void closeAsync(AutoCloseable stream) { + Schedulers.io().scheduleDirect(() -> closeQuietly(stream)); + } + + private static void closeQuietly(AutoCloseable stream) { + try { + stream.close(); + } catch (Exception e) { + logger.debug("Failed to close the response stream", e); + } + } + private static LlmResponse withCacheMetadata( LlmResponse response, @Nullable CacheMetadata cacheMetadata) { return cacheMetadata == null diff --git a/core/src/test/java/com/google/adk/flows/llmflows/BaseLlmFlowCancelTest.java b/core/src/test/java/com/google/adk/flows/llmflows/BaseLlmFlowCancelTest.java new file mode 100644 index 000000000..c7f34e717 --- /dev/null +++ b/core/src/test/java/com/google/adk/flows/llmflows/BaseLlmFlowCancelTest.java @@ -0,0 +1,303 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.adk.flows.llmflows; + +import static com.google.adk.testing.TestUtils.createInvocationContext; +import static com.google.adk.testing.TestUtils.createLlmResponse; +import static com.google.adk.testing.TestUtils.createTestAgentBuilder; +import static com.google.adk.testing.TestUtils.createTestLlm; +import static com.google.common.truth.Truth.assertThat; +import static java.util.concurrent.TimeUnit.SECONDS; + +import com.google.adk.agents.BaseAgent; +import com.google.adk.agents.InvocationContext; +import com.google.adk.agents.LlmAgent; +import com.google.adk.agents.ParallelAgent; +import com.google.adk.events.Event; +import com.google.adk.events.EventActions; +import com.google.adk.models.LlmResponse; +import com.google.adk.runner.InMemoryRunner; +import com.google.adk.sessions.Session; +import com.google.adk.testing.TestLlm; +import com.google.adk.tools.BaseTool; +import com.google.adk.tools.ToolContext; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.genai.types.Content; +import com.google.genai.types.FunctionDeclaration; +import com.google.genai.types.Part; +import io.reactivex.rxjava3.core.Flowable; +import io.reactivex.rxjava3.core.Single; +import io.reactivex.rxjava3.disposables.Disposable; +import io.reactivex.rxjava3.processors.PublishProcessor; +import io.reactivex.rxjava3.schedulers.TestScheduler; +import io.reactivex.rxjava3.subscribers.DisposableSubscriber; +import io.reactivex.rxjava3.subscribers.TestSubscriber; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.CountDownLatch; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests that cancelling a {@link BaseLlmFlow#run} subscription reaches the model stream. */ +@RunWith(JUnit4.class) +public final class BaseLlmFlowCancelTest { + + private static final Content FUNCTION_CALL = + Content.fromParts(Part.fromFunctionCall("my_function", ImmutableMap.of("arg1", "value1"))); + private static final Content TEXT = Content.fromParts(Part.fromText("LLM response")); + + @Test + public void run_disposedWhileModelStreams_cancelsModelStream() throws Exception { + NeverEndingModel model = new NeverEndingModel(); + InvocationContext invocationContext = + createInvocationContext(createTestAgentBuilder(model.llm).build()); + + Disposable subscription = + createBaseLlmFlow().run(invocationContext).subscribe(event -> {}, error -> {}); + assertThat(model.subscribed.await(5, SECONDS)).isTrue(); + subscription.dispose(); + + assertThat(model.cancelled.await(2, SECONDS)).isTrue(); + } + + @Test + public void runner_disposedWhileModelStreams_cancelsModelStream() throws Exception { + NeverEndingModel model = new NeverEndingModel(); + LlmAgent agent = createTestAgentBuilder(model.llm).build(); + InMemoryRunner runner = new InMemoryRunner(agent); + Session session = runner.sessionService().createSession(runner.appName(), "user").blockingGet(); + + Disposable subscription = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("hi"))) + .subscribe(event -> {}, error -> {}); + assertThat(model.subscribed.await(5, SECONDS)).isTrue(); + subscription.dispose(); + + assertThat(model.cancelled.await(2, SECONDS)).isTrue(); + } + + @Test + public void run_disposedBeforeFunctionCallArrives_doesNotRunTool() throws Exception { + PublishProcessor responses = PublishProcessor.create(); + CountDownLatch cancelled = new CountDownLatch(1); + TestLlm testLlm = createTestLlm(() -> responses.doOnCancel(cancelled::countDown)); + RecordingTool tool = new RecordingTool(); + InvocationContext invocationContext = + createInvocationContext( + createTestAgentBuilder(testLlm).tools(ImmutableList.of(tool)).build()); + + Disposable subscription = + createBaseLlmFlow().run(invocationContext).subscribe(event -> {}, error -> {}); + assertThat(responses.hasSubscribers()).isTrue(); + subscription.dispose(); + assertThat(cancelled.await(2, SECONDS)).isTrue(); + responses.onNext(createLlmResponse(FUNCTION_CALL)); + responses.onComplete(); + + assertThat(tool.ran.getCount()).isEqualTo(1); + assertThat(testLlm.getRequests()).hasSize(1); + } + + @Test + public void run_consumedToCompletion_runsToolAndEmitsEveryEvent() throws Exception { + TestLlm testLlm = + createTestLlm( + Flowable.just(createLlmResponse(FUNCTION_CALL)), + Flowable.just(createLlmResponse(TEXT))); + RecordingTool tool = new RecordingTool(); + InvocationContext invocationContext = + createInvocationContext( + createTestAgentBuilder(testLlm).tools(ImmutableList.of(tool)).build()); + + List events = createBaseLlmFlow().run(invocationContext).toList().blockingGet(); + + assertThat(tool.ran.getCount()).isEqualTo(0); + assertThat(events).hasSize(3); + assertThat(events.get(2).content()).hasValue(TEXT); + assertThat(testLlm.getRequests()).hasSize(2); + } + + @Test + public void run_boundedDemand_deliversFinalEventAndCompletion() { + TestLlm testLlm = createTestLlm(createLlmResponse(TEXT)); + InvocationContext invocationContext = + createInvocationContext(createTestAgentBuilder(testLlm).build()); + + TestSubscriber subscriber = createBaseLlmFlow().run(invocationContext).test(1); + + subscriber.assertValueCount(1).assertComplete(); + assertThat(testLlm.getRequests()).hasSize(1); + } + + @Test + public void run_boundedDemand_deliversModelErrorAfterBufferedEvent() { + TestLlm testLlm = + createTestLlm( + () -> + Flowable.concat( + Flowable.just(createLlmResponse(TEXT)), + Flowable.error(new IllegalStateException("model failed")))); + InvocationContext invocationContext = + createInvocationContext(createTestAgentBuilder(testLlm).build()); + + TestSubscriber subscriber = createBaseLlmFlow().run(invocationContext).test(1); + + subscriber.assertValueCount(1).assertError(IllegalStateException.class); + } + + @Test + public void run_cancelledInOnSubscribe_doesNotLeaveModelStreaming() throws Exception { + NeverEndingModel model = new NeverEndingModel(); + InvocationContext invocationContext = + createInvocationContext(createTestAgentBuilder(model.llm).build()); + // A single step subscribes to the cached step directly, with no concatWith in front of it. + BaseLlmFlow singleStepFlow = + new BaseLlmFlow(ImmutableList.of(), ImmutableList.of(), Optional.of(1)) {}; + + singleStepFlow + .run(invocationContext) + .subscribe( + new DisposableSubscriber() { + @Override + protected void onStart() { + cancel(); + } + + @Override + public void onNext(Event event) {} + + @Override + public void onError(Throwable error) {} + + @Override + public void onComplete() {} + }); + + boolean neverSubscribed = model.subscribed.getCount() == 1; + assertThat(neverSubscribed || model.cancelled.await(2, SECONDS)).isTrue(); + } + + @Test + public void run_subscribedAgainAfterCancel_startsNewRunAndCompletes() throws Exception { + PublishProcessor responses = PublishProcessor.create(); + CountDownLatch cancelled = new CountDownLatch(1); + TestLlm testLlm = createTestLlm(() -> responses.doOnCancel(cancelled::countDown)); + InvocationContext invocationContext = + createInvocationContext(createTestAgentBuilder(testLlm).build()); + Flowable events = createBaseLlmFlow().run(invocationContext); + + TestSubscriber first = events.test(); + responses.onNext(createLlmResponse(TEXT)); + first.assertValueCount(1).assertNotComplete(); + first.cancel(); + assertThat(cancelled.await(2, SECONDS)).isTrue(); + + TestSubscriber second = events.test(); + responses.onNext(createLlmResponse(TEXT)); + responses.onComplete(); + + second.assertValueCount(1).assertComplete(); + assertThat(second.values().get(0).content()).hasValue(TEXT); + assertThat(testLlm.getRequests()).hasSize(2); + } + + @Test + public void parallelAgent_subAgentEscalates_cancelsSiblingModelStream() throws Exception { + NeverEndingModel model = new NeverEndingModel(); + LlmAgent streamingAgent = createTestAgentBuilder(model.llm).name("streaming_agent").build(); + TestScheduler scheduler = new TestScheduler(); + ParallelAgent parallelAgent = + ParallelAgent.builder() + .name("parallel_agent") + .subAgents(streamingAgent, new EscalatingAgent()) + .scheduler(scheduler) + .build(); + InvocationContext invocationContext = createInvocationContext(parallelAgent); + + TestSubscriber subscriber = parallelAgent.runAsync(invocationContext).test(); + // Subscribes the branches in order: the model starts streaming, then the sibling escalates. + scheduler.triggerActions(); + + assertThat(model.subscribed.getCount()).isEqualTo(0); + subscriber.assertValueCount(1).assertComplete(); + assertThat(subscriber.values().get(0).actions().escalate()).hasValue(true); + assertThat(model.cancelled.await(2, SECONDS)).isTrue(); + } + + private static BaseLlmFlow createBaseLlmFlow() { + return new BaseLlmFlow(ImmutableList.of(), ImmutableList.of(), Optional.empty()) {}; + } + + /** A model whose stream never completes and reports when it is subscribed to and cancelled. */ + private static final class NeverEndingModel { + final CountDownLatch subscribed = new CountDownLatch(1); + final CountDownLatch cancelled = new CountDownLatch(1); + final TestLlm llm = + createTestLlm( + () -> + Flowable.never() + .doOnSubscribe(subscription -> subscribed.countDown()) + .doOnCancel(cancelled::countDown)); + } + + /** A sub-agent that escalates with its first event. */ + private static final class EscalatingAgent extends BaseAgent { + EscalatingAgent() { + super("escalating_agent", "escalates immediately", ImmutableList.of(), null, null); + } + + @Override + protected Flowable runAsyncImpl(InvocationContext invocationContext) { + return Flowable.just( + Event.builder() + .author(name()) + .invocationId(invocationContext.invocationId()) + .branch(invocationContext.branch().orElse(null)) + .content(Content.fromParts(Part.fromText("escalating"))) + .actions(EventActions.builder().escalate(true).build()) + .build()); + } + + @Override + protected Flowable runLiveImpl(InvocationContext invocationContext) { + throw new UnsupportedOperationException("Not implemented"); + } + } + + private static final class RecordingTool extends BaseTool { + final CountDownLatch ran = new CountDownLatch(1); + + RecordingTool() { + super("my_function", "tool description for my_function"); + } + + @Override + public Optional declaration() { + return Optional.of(FunctionDeclaration.builder().name(name()).build()); + } + + @Override + public Single> runAsync(Map args, ToolContext toolContext) { + ran.countDown(); + return Single.just(ImmutableMap.of("response", "response for my_function")); + } + } +} diff --git a/core/src/test/java/com/google/adk/models/GeminiContinuationTest.java b/core/src/test/java/com/google/adk/models/GeminiContinuationTest.java index f1b9cd5eb..42f2575fe 100644 --- a/core/src/test/java/com/google/adk/models/GeminiContinuationTest.java +++ b/core/src/test/java/com/google/adk/models/GeminiContinuationTest.java @@ -21,6 +21,7 @@ import static com.google.common.collect.Iterables.getOnlyElement; import static com.google.common.truth.Truth.assertThat; import static java.nio.charset.StandardCharsets.UTF_8; +import static java.util.concurrent.TimeUnit.SECONDS; import static java.util.stream.Collectors.joining; import com.fasterxml.jackson.databind.JsonNode; @@ -28,20 +29,27 @@ import com.google.adk.agents.ContextCacheConfig; import com.google.common.collect.ImmutableList; import com.google.genai.Client; +import com.google.genai.ResponseStream; import com.google.genai.types.Candidate; import com.google.genai.types.ClientOptions; import com.google.genai.types.Content; import com.google.genai.types.FinishReason; +import com.google.genai.types.GenerateContentConfig; import com.google.genai.types.GenerateContentResponse; import com.google.genai.types.GenerateContentResponseUsageMetadata; import com.google.genai.types.Part; +import io.reactivex.rxjava3.schedulers.Schedulers; +import io.reactivex.rxjava3.subscribers.TestSubscriber; import java.io.IOException; import java.util.ArrayDeque; import java.util.ArrayList; import java.util.Arrays; import java.util.Base64; import java.util.Deque; +import java.util.Iterator; import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; import okhttp3.Interceptor; import okhttp3.MediaType; import okhttp3.OkHttpClient; @@ -50,6 +58,9 @@ import okhttp3.Response; import okhttp3.ResponseBody; import okio.Buffer; +import okio.BufferedSource; +import okio.ForwardingSource; +import okio.Okio; import org.jspecify.annotations.Nullable; import org.junit.Test; import org.junit.runner.RunWith; @@ -184,6 +195,41 @@ public void generateContent_pausedStreamWithCache_isResumedWithTheCache() { assertUsage(last, /* prompt= */ 10 + 15, /* candidates= */ 5 + 3, /* total= */ 15 + 18); } + @Test + public void readStream_cancelWhileResumeOpens_closesBothStreams() throws Exception { + api.respondStream( + response("CONTINUATION", "state", usage(10, 5, 15), Part.fromText("The answer is"))); + api.respondStream(response("STOP", /* token= */ null, usage(15, 3, 18), Part.fromText(" 42."))); + CountDownLatch resumeRequested = new CountDownLatch(1); + CountDownLatch releaseResume = new CountDownLatch(1); + api.holdRequest(/* index= */ 1, resumeRequested, releaseResume); + // Signals after the SDK's close() returns, so the resume is released only once the first + // close can no longer see it. + CountDownLatch firstCloseReturned = new CountDownLatch(1); + CompletableFuture future = + api.client() + .async + .models + .generateContentStream( + "gemini-test-model", + ImmutableList.of(QUESTION), + GenerateContentConfig.builder().build()) + .thenApply(stream -> new CloseTrackingStream(stream, firstCloseReturned)); + + TestSubscriber subscriber = + Gemini.readStream(future).subscribeOn(Schedulers.io()).test(); + try { + assertThat(resumeRequested.await(5, SECONDS)).isTrue(); + subscriber.cancel(); + assertThat(api.bodyClosed(0).await(5, SECONDS)).isTrue(); + assertThat(firstCloseReturned.await(5, SECONDS)).isTrue(); + } finally { + releaseResume.countDown(); + } + + assertThat(api.bodyClosed(1).await(5, SECONDS)).isTrue(); + } + private List generate(boolean stream) { LlmRequest request = LlmRequest.builder().contents(ImmutableList.of(QUESTION)).build(); return gemini.generateContent(request, stream).toList().blockingGet(); @@ -273,10 +319,38 @@ private static String base64(String token) { return Base64.getEncoder().encodeToString(token.getBytes(UTF_8)); } + /** Wraps an SDK stream and counts down a latch once its {@code close()} returns. */ + private static final class CloseTrackingStream + implements Iterable, AutoCloseable { + private final ResponseStream delegate; + private final CountDownLatch closeReturned; + + CloseTrackingStream( + ResponseStream delegate, CountDownLatch closeReturned) { + this.delegate = delegate; + this.closeReturned = closeReturned; + } + + @Override + public Iterator iterator() { + return delegate.iterator(); + } + + @Override + public void close() { + delegate.close(); + closeReturned.countDown(); + } + } + /** A fake Gemini API that records each request and answers it with the next canned response. */ private static final class ScriptedGeminiApi implements Interceptor { private final Deque responses = new ArrayDeque<>(); private final List requests = new ArrayList<>(); + private final List bodiesClosed = new ArrayList<>(); + private int heldIndex = -1; + private CountDownLatch heldRequested; + private CountDownLatch heldRelease; Client client() { OkHttpClient httpClient = new OkHttpClient.Builder().addInterceptor(this).build(); @@ -306,22 +380,70 @@ synchronized void respondWithCache(String name) { responses.add("{\"name\":\"" + name + "\",\"expireTime\":\"2099-01-01T00:00:00Z\"}"); } + /** + * Makes request {@code index} (0-based) count down {@code requested} when it arrives and wait + * for {@code release} before it is answered. + */ + synchronized void holdRequest(int index, CountDownLatch requested, CountDownLatch release) { + heldIndex = index; + heldRequested = requested; + heldRelease = release; + } + + /** Returns a latch that opens when the body of request {@code index} (0-based) is closed. */ + synchronized CountDownLatch bodyClosed(int index) { + while (bodiesClosed.size() <= index) { + bodiesClosed.add(new CountDownLatch(1)); + } + return bodiesClosed.get(index); + } + synchronized ImmutableList requests() { return ImmutableList.copyOf(requests); } @Override - public synchronized Response intercept(Chain chain) throws IOException { + public Response intercept(Chain chain) throws IOException { Request request = chain.request(); Buffer body = new Buffer(); request.body().writeTo(body); - requests.add(objectMapper.readTree(body.readUtf8())); + int index; + String answer; + CountDownLatch closed; + CountDownLatch release = null; + synchronized (this) { + index = requests.size(); + requests.add(objectMapper.readTree(body.readUtf8())); + answer = responses.remove(); + closed = bodyClosed(index); + if (index == heldIndex) { + heldRequested.countDown(); + release = heldRelease; + } + } + if (release != null) { + try { + release.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IOException(e); + } + } + BufferedSource source = + Okio.buffer( + new ForwardingSource(new Buffer().writeUtf8(answer)) { + @Override + public void close() throws IOException { + closed.countDown(); + super.close(); + } + }); return new Response.Builder() .request(request) .protocol(Protocol.HTTP_1_1) .code(200) .message("fake") - .body(ResponseBody.create(responses.remove(), MediaType.get("application/json"))) + .body(ResponseBody.create(source, MediaType.get("application/json"), -1)) .build(); } } diff --git a/core/src/test/java/com/google/adk/models/GeminiTest.java b/core/src/test/java/com/google/adk/models/GeminiTest.java index a56628493..bc5c02668 100644 --- a/core/src/test/java/com/google/adk/models/GeminiTest.java +++ b/core/src/test/java/com/google/adk/models/GeminiTest.java @@ -17,6 +17,7 @@ import static com.google.common.truth.Truth.assertThat; import static java.nio.charset.StandardCharsets.UTF_8; +import static java.util.concurrent.TimeUnit.SECONDS; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; @@ -34,7 +35,13 @@ import com.google.genai.types.ToolResponse; import io.reactivex.rxjava3.core.Flowable; import io.reactivex.rxjava3.functions.Predicate; +import io.reactivex.rxjava3.schedulers.Schedulers; import io.reactivex.rxjava3.subscribers.TestSubscriber; +import java.util.Iterator; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.stream.IntStream; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -2192,4 +2199,72 @@ private static GenerateContentResponseUsageMetadata createUsageMetadata( .totalTokenCount(totalTokens) .build(); } + + @Test + public void readStream_cancelMidStream_closesStream() throws Exception { + FakeStream stream = new FakeStream(); + + TestSubscriber subscriber = + Gemini.readStream(CompletableFuture.completedFuture(stream)).take(1).test(); + + subscriber.assertValues("chunk-1").assertComplete(); + assertThat(stream.closed.await(5, SECONDS)).isTrue(); + } + + @Test + public void readStream_cancelBeforeStreamOpens_closesStreamOnceOpen() throws Exception { + CountDownLatch waiting = new CountDownLatch(1); + CompletableFuture future = + new CompletableFuture<>() { + @Override + public FakeStream get() throws InterruptedException, ExecutionException { + waiting.countDown(); + return super.get(); + } + }; + TestSubscriber subscriber = + Gemini.readStream(future).subscribeOn(Schedulers.io()).test(); + assertThat(waiting.await(5, SECONDS)).isTrue(); + + subscriber.cancel(); + FakeStream stream = new FakeStream(); + future.complete(stream); + + assertThat(stream.closed.await(5, SECONDS)).isTrue(); + subscriber.assertNoValues(); + } + + @Test + public void readStream_readsToEnd_emitsEveryResponse() { + FakeStream stream = new FakeStream(3); + + Gemini.readStream(CompletableFuture.completedFuture(stream)) + .test() + .assertValues("chunk-1", "chunk-2", "chunk-3") + .assertComplete(); + } + + /** A stream of "chunk-N" strings that records when it is closed; unbounded by default. */ + private static final class FakeStream implements Iterable, AutoCloseable { + final CountDownLatch closed = new CountDownLatch(1); + private final int size; + + FakeStream() { + this(Integer.MAX_VALUE); + } + + FakeStream(int size) { + this.size = size; + } + + @Override + public Iterator iterator() { + return IntStream.rangeClosed(1, size).mapToObj(i -> "chunk-" + i).iterator(); + } + + @Override + public void close() { + closed.countDown(); + } + } }