fix: close the Gemini response stream after a read error

PiperOrigin-RevId: 1008122641

Author: google-genai-bot

core/src/main/java/com/google/adk/models/Gemini.java

@@ -330,26 +330,26 @@ private static GenerateContentConfig withTrackingHeaders(@Nullable GenerateConte
   }
 
   /**
-   * 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.
+   * Reads the stream and closes it on {@link Schedulers#io()} unless it was read to the end, so a
+   * cancel or a read error 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 <T, S extends Iterable<T> & AutoCloseable> Flowable<T> readStream(
       CompletableFuture<S> streamFuture) {
     return Flowable.defer(
         () -> {
           AtomicBoolean cancelled = new AtomicBoolean();
-          AtomicBoolean terminated = new AtomicBoolean();
+          AtomicBoolean completed = new AtomicBoolean();
           return Flowable.using(
               () -> streamFuture,
               future ->
                   Flowable.fromFuture(future)
                       .flatMapIterable(stream -> closingAfterCancel(stream, cancelled))
-                      .doOnTerminate(() -> terminated.set(true)),
+                      .doOnComplete(() -> completed.set(true)),
               future -> {
-                if (!terminated.get()) {
+                if (!completed.get()) {
                   cancelled.set(true);
                   var unused = future.thenAccept(Gemini::closeAsync);
                 }

core/src/test/java/com/google/adk/flows/llmflows/BaseLlmFlowCancelTest.java

@@ -51,6 +51,7 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.atomic.AtomicInteger;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -112,7 +113,7 @@ public void run_disposedBeforeFunctionCallArrives_doesNotRunTool() throws Except
     responses.onNext(createLlmResponse(FUNCTION_CALL));
     responses.onComplete();
 
-    assertThat(tool.ran.getCount()).isEqualTo(1);
+    assertThat(tool.runCount.get()).isEqualTo(0);
     assertThat(testLlm.getRequests()).hasSize(1);
   }
 
@@ -129,7 +130,7 @@ public void run_consumedToCompletion_runsToolAndEmitsEveryEvent() throws Excepti
 
     List<Event> events = createBaseLlmFlow().run(invocationContext).toList().blockingGet();
 
-    assertThat(tool.ran.getCount()).isEqualTo(0);
+    assertThat(tool.runCount.get()).isEqualTo(1);
     assertThat(events).hasSize(3);
     assertThat(events.get(2).content()).hasValue(TEXT);
     assertThat(testLlm.getRequests()).hasSize(2);
@@ -283,7 +284,7 @@ protected Flowable<Event> runLiveImpl(InvocationContext invocationContext) {
   }
 
   private static final class RecordingTool extends BaseTool {
-    final CountDownLatch ran = new CountDownLatch(1);
+    final AtomicInteger runCount = new AtomicInteger();
 
     RecordingTool() {
       super("my_function", "tool description for my_function");
@@ -296,7 +297,7 @@ public Optional<FunctionDeclaration> declaration() {
 
     @Override
     public Single<Map<String, Object>> runAsync(Map<String, Object> args, ToolContext toolContext) {
-      ran.countDown();
+      runCount.incrementAndGet();
       return Single.just(ImmutableMap.of("response", "response for my_function"));
     }
   }

core/src/test/java/com/google/adk/models/GeminiTest.java

@@ -2234,6 +2234,18 @@ public FakeStream get() throws InterruptedException, ExecutionException {
     subscriber.assertNoValues();
   }
 
+  @Test
+  public void readStream_errorMidStream_closesStream() throws Exception {
+    IllegalStateException failure = new IllegalStateException("stream failed");
+    FakeStream stream = new FailingStream(2, failure);
+
+    Gemini.readStream(CompletableFuture.completedFuture(stream))
+        .test()
+        .assertValues("chunk-1", "chunk-2")
+        .assertError(failure);
+    assertThat(stream.closed.await(5, SECONDS)).isTrue();
+  }
+
   @Test
   public void readStream_readsToEnd_emitsEveryResponse() {
     FakeStream stream = new FakeStream(3);
@@ -2245,7 +2257,7 @@ public void readStream_readsToEnd_emitsEveryResponse() {
   }
 
   /** A stream of "chunk-N" strings that records when it is closed; unbounded by default. */
-  private static final class FakeStream implements Iterable<String>, AutoCloseable {
+  private static class FakeStream implements Iterable<String>, AutoCloseable {
     final CountDownLatch closed = new CountDownLatch(1);
     private final int size;
 
@@ -2267,4 +2279,33 @@ public void close() {
       closed.countDown();
     }
   }
+
+  /** A {@link FakeStream} of {@code size} chunks whose iterator then throws {@code failure}. */
+  private static final class FailingStream extends FakeStream {
+    private final RuntimeException failure;
+
+    FailingStream(int size, RuntimeException failure) {
+      super(size);
+      this.failure = failure;
+    }
+
+    @Override
+    public Iterator<String> iterator() {
+      Iterator<String> chunks = super.iterator();
+      return new Iterator<String>() {
+        @Override
+        public boolean hasNext() {
+          if (!chunks.hasNext()) {
+            throw failure;
+          }
+          return true;
+        }
+
+        @Override
+        public String next() {
+          return chunks.next();
+        }
+      };
+    }
+  }
 }