Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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.
*/
Comment thread
MiloszSobczyk marked this conversation as resolved.
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.
*
Expand Down
81 changes: 76 additions & 5 deletions core/src/main/java/com/google/adk/models/Gemini.java
Original file line number Diff line number Diff line change
Expand Up @@ -38,14 +38,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;
Expand Down Expand Up @@ -299,11 +302,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(
Expand All @@ -316,6 +315,78 @@ private Flowable<LlmResponse> sendRequest(
}
}

/**
* 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.
*/
Comment thread
MiloszSobczyk marked this conversation as resolved.
@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);
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
Expand Down
Loading
Loading