Skip to content

Commit 216bec2

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 48e9ae7 commit 216bec2

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
@@ -35,14 +35,17 @@
3535
import com.google.genai.types.Part;
3636
import com.google.genai.types.PartialArg;
3737
import io.reactivex.rxjava3.core.Flowable;
38+
import io.reactivex.rxjava3.schedulers.Schedulers;
3839
import java.time.Duration;
3940
import java.util.ArrayList;
41+
import java.util.Iterator;
4042
import java.util.LinkedHashMap;
4143
import java.util.List;
4244
import java.util.Map;
4345
import java.util.Objects;
4446
import java.util.concurrent.CompletableFuture;
4547
import java.util.concurrent.ExecutorService;
48+
import java.util.concurrent.atomic.AtomicBoolean;
4649
import okhttp3.OkHttpClient;
4750
import org.jspecify.annotations.Nullable;
4851
import org.slf4j.Logger;
@@ -296,11 +299,7 @@ private Flowable<LlmResponse> sendRequest(
296299
effectiveModelName, llmRequest.contents(), config);
297300

298301
return Flowable.defer(
299-
() ->
300-
new StreamingResponseAggregator()
301-
.process(
302-
Flowable.fromFuture(streamFuture).flatMapIterable(iterable -> iterable),
303-
cacheMetadata));
302+
() -> new StreamingResponseAggregator().process(readStream(streamFuture), cacheMetadata));
304303
} else {
305304
logger.debug("Sending generateContent request to model {}", effectiveModelName);
306305
return Flowable.fromFuture(
@@ -330,6 +329,78 @@ private static GenerateContentConfig withTrackingHeaders(@Nullable GenerateConte
330329
.build();
331330
}
332331

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

0 commit comments

Comments
 (0)