Skip to content

Commit dd914fc

Browse files
committed
fix(flows): propagate downstream cancellation to the model stream
- Cancelling the `Flowable` from `Runner.runAsync` now cancels the subscription to the current LLM step. `Gemini` schedules a close of the streamed response on cancel, including a response that opens after the cancel. Cancelling an active step also completes its shared cache; later subscribers to that cache replay buffered events and completion. - Active tool subscriptions are disposed on cancel. With two or more calls in `ToolExecutionMode.PARALLEL_SUBSCRIBE`, tool subscriptions use the agent's executor when set, otherwise the invocation's scheduler (`Schedulers.io()` by default). Disposing a running subscription task on the default IO scheduler from another thread interrupts its worker; other schedulers may not. - An error that becomes undeliverable after a cancel can reach `RxJavaPlugins.onError`, usually wrapped in `UndeliverableException`. Without a global handler RxJava passes it to the thread's uncaught exception handler, which can terminate an Android app. - Escalation in `ParallelAgent` now cancels the sibling sub-agents' active LLM steps and tools. Escalation in `LoopAgent`, outside durable resumable mode, now cancels the remaining LLM step and tools of the active sub-agent.
1 parent 8bc04da commit dd914fc

5 files changed

Lines changed: 590 additions & 9 deletions

File tree

‎core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@
5555
import io.reactivex.rxjava3.core.Single;
5656
import io.reactivex.rxjava3.disposables.Disposable;
5757
import io.reactivex.rxjava3.observers.DisposableCompletableObserver;
58+
import io.reactivex.rxjava3.subjects.CompletableSubject;
5859
import java.util.ArrayList;
5960
import java.util.List;
6061
import java.util.Optional;
@@ -520,7 +521,8 @@ public Flowable<Event> run(InvocationContext invocationContext) {
520521

521522
private Flowable<Event> run(
522523
Context spanContext, InvocationContext invocationContext, int stepsCompleted) {
523-
Flowable<Event> currentStepEvents = runOneStep(spanContext, invocationContext).cache();
524+
Flowable<Event> currentStepEvents =
525+
cancellableCache(runOneStep(spanContext, invocationContext));
524526
if (stepsCompleted + 1 >= maxSteps) {
525527
logger.debug("Ending flow execution because max steps reached.");
526528
return currentStepEvents;
@@ -554,6 +556,16 @@ private Flowable<Event> run(
554556
}));
555557
}
556558

559+
/**
560+
* Like {@code step.cache()}, but cancels {@code step} when a subscriber cancels. {@code cache()}
561+
* never cancels upstream, so a client disconnect left the model stream and pending tool calls
562+
* running. A subscriber that arrives after the cancel still gets the buffered events.
563+
*/
564+
private static Flowable<Event> cancellableCache(Flowable<Event> step) {
565+
CompletableSubject cancelled = CompletableSubject.create();
566+
return step.takeUntil(cancelled.toFlowable()).cache().doOnCancel(cancelled::onComplete);
567+
}
568+
557569
/**
558570
* Executes the LLM flow in streaming mode.
559571
*

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

Lines changed: 76 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -38,14 +38,17 @@
3838
import com.google.genai.types.Part;
3939
import com.google.genai.types.PartialArg;
4040
import io.reactivex.rxjava3.core.Flowable;
41+
import io.reactivex.rxjava3.schedulers.Schedulers;
4142
import java.time.Duration;
4243
import java.util.ArrayList;
44+
import java.util.Iterator;
4345
import java.util.LinkedHashMap;
4446
import java.util.List;
4547
import java.util.Map;
4648
import java.util.Objects;
4749
import java.util.concurrent.CompletableFuture;
4850
import java.util.concurrent.ExecutorService;
51+
import java.util.concurrent.atomic.AtomicBoolean;
4952
import okhttp3.OkHttpClient;
5053
import org.jspecify.annotations.Nullable;
5154
import org.slf4j.Logger;
@@ -299,11 +302,7 @@ private Flowable<LlmResponse> sendRequest(
299302
effectiveModelName, llmRequest.contents(), config);
300303

301304
return Flowable.defer(
302-
() ->
303-
new StreamingResponseAggregator()
304-
.process(
305-
Flowable.fromFuture(streamFuture).flatMapIterable(iterable -> iterable),
306-
cacheMetadata));
305+
() -> new StreamingResponseAggregator().process(readStream(streamFuture), cacheMetadata));
307306
} else {
308307
logger.debug("Sending generateContent request to model {}", effectiveModelName);
309308
return Flowable.fromFuture(
@@ -316,6 +315,78 @@ private Flowable<LlmResponse> sendRequest(
316315
}
317316
}
318317

318+
/**
319+
* Reads the stream and, when the subscriber cancels, closes it on {@link Schedulers#io()} so a
320+
* partial read doesn't leave the HTTP response open. The close waits for the stream to open and
321+
* can block behind a {@code readLine()} in progress. An iterator call that returns after the
322+
* cancel closes the stream again, covering a continuation request opened during that call.
323+
*/
324+
@VisibleForTesting
325+
static <T, S extends Iterable<T> & AutoCloseable> Flowable<T> readStream(
326+
CompletableFuture<S> streamFuture) {
327+
return Flowable.defer(
328+
() -> {
329+
AtomicBoolean cancelled = new AtomicBoolean();
330+
AtomicBoolean terminated = new AtomicBoolean();
331+
return Flowable.using(
332+
() -> streamFuture,
333+
future ->
334+
Flowable.fromFuture(future)
335+
.flatMapIterable(stream -> closingAfterCancel(stream, cancelled))
336+
.doOnTerminate(() -> terminated.set(true)),
337+
future -> {
338+
if (!terminated.get()) {
339+
cancelled.set(true);
340+
future.thenAccept(Gemini::closeAsync);
341+
}
342+
},
343+
/* eager= */ false);
344+
});
345+
}
346+
347+
private static <T, S extends Iterable<T> & AutoCloseable> Iterable<T> closingAfterCancel(
348+
S stream, AtomicBoolean cancelled) {
349+
Iterator<T> delegate = stream.iterator();
350+
return () ->
351+
new Iterator<T>() {
352+
@Override
353+
public boolean hasNext() {
354+
try {
355+
return delegate.hasNext();
356+
} finally {
357+
closeIfCancelled();
358+
}
359+
}
360+
361+
@Override
362+
public T next() {
363+
try {
364+
return delegate.next();
365+
} finally {
366+
closeIfCancelled();
367+
}
368+
}
369+
370+
private void closeIfCancelled() {
371+
if (cancelled.get()) {
372+
closeAsync(stream);
373+
}
374+
}
375+
};
376+
}
377+
378+
private static void closeAsync(AutoCloseable stream) {
379+
Schedulers.io().scheduleDirect(() -> closeQuietly(stream));
380+
}
381+
382+
private static void closeQuietly(AutoCloseable stream) {
383+
try {
384+
stream.close();
385+
} catch (Exception e) {
386+
logger.debug("Failed to close the response stream", e);
387+
}
388+
}
389+
319390
private static LlmResponse withCacheMetadata(
320391
LlmResponse response, @Nullable CacheMetadata cacheMetadata) {
321392
return cacheMetadata == null

0 commit comments

Comments
 (0)