Merge pull request #1621 from sangkyoonnam:fix/1618-cancel-propagation

PiperOrigin-RevId: 1007989909

Author: copybara-github

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<Event> runStepAndCheckContinue(
       InvocationContext invocationContext,
       int stepsCompleted,
       AtomicBoolean continueFlow) {
-    Flowable<Event> currentStepEvents = runOneStep(spanContext, invocationContext).cache();
+    Flowable<Event> 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<Event> 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<Event> cancellableCache(Flowable<Event> step) {
+    CompletableSubject cancelled = CompletableSubject.create();
+    return step.takeUntil(cancelled.toFlowable()).cache().doOnCancel(cancelled::onComplete);
+  }
+
   /**
    * Executes the LLM flow in streaming mode.
    *

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

@@ -35,14 +35,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;
@@ -296,11 +299,7 @@ private Flowable<LlmResponse> 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(
@@ -330,6 +329,78 @@ private static GenerateContentConfig withTrackingHeaders(@Nullable GenerateConte
         .build();
   }
 
+  /**
+   * 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 <T, S extends Iterable<T> & AutoCloseable> Flowable<T> readStream(
+      CompletableFuture<S> 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);
+                  var unused = future.thenAccept(Gemini::closeAsync);
+                }
+              },
+              /* eager= */ false);
+        });
+  }
+
+  private static <T, S extends Iterable<T> & AutoCloseable> Iterable<T> closingAfterCancel(
+      S stream, AtomicBoolean cancelled) {
+    Iterator<T> delegate = stream.iterator();
+    return () ->
+        new Iterator<T>() {
+          @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