Skip to content

Commit 542ad25

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 0ee0d4b commit 542ad25

5 files changed

Lines changed: 592 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;
@@ -542,7 +543,8 @@ private Flowable<Event> runStepAndCheckContinue(
542543
InvocationContext invocationContext,
543544
int stepsCompleted,
544545
AtomicBoolean continueFlow) {
545-
Flowable<Event> currentStepEvents = runOneStep(spanContext, invocationContext).cache();
546+
Flowable<Event> currentStepEvents =
547+
cancellableCache(runOneStep(spanContext, invocationContext));
546548
if (stepsCompleted + 1 >= maxSteps) {
547549
logger.debug("Ending flow execution because max steps reached.");
548550
return currentStepEvents;
@@ -576,6 +578,16 @@ private Flowable<Event> runStepAndCheckContinue(
576578
}));
577579
}
578580

581+
/**
582+
* Like {@code step.cache()}, but cancels {@code step} when a subscriber cancels. {@code cache()}
583+
* never cancels upstream, so a client disconnect left the model stream and pending tool calls
584+
* running. A subscriber that arrives after the cancel still gets the buffered events.
585+
*/
586+
private static Flowable<Event> cancellableCache(Flowable<Event> step) {
587+
CompletableSubject cancelled = CompletableSubject.create();
588+
return step.takeUntil(cancelled.toFlowable()).cache().doOnCancel(cancelled::onComplete);
589+
}
590+
579591
/**
580592
* Executes the LLM flow in streaming mode.
581593
*

‎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)