Skip to content

Commit 9a98d4d

Browse files
committed
fix(flows): avoid StackOverflowError after many steps in BaseLlmFlow
BaseLlmFlow.run subscribed to the next step from inside the previous step's completion, through concatWith, toList, flatMapPublisher and PersistBarrier.awaitPersisted(...).andThen(run(...)). When a step completed on the subscribing thread, as it does with the core models and the session services, each step ran 23 stack frames deeper than the one before. With a 1 MB thread stack, the default on x86-64, an agent that kept calling tools failed with StackOverflowError after roughly 250-300 LLM calls, before the default maxLlmCalls of 500 could end the run, and the error never reached onError. Drive the steps with repeatUntil instead, as LoopAgent does on its resumable path. repeatUntil resubscribes from a loop when the source completes synchronously, so every step starts at the same stack depth, like the while loop in adk-python's BaseLlmFlow.run_async. The end-of-flow checks, the pause on a pending long-running call, the PersistBarrier wait between steps and maxSteps are unchanged.
1 parent 2d4a59e commit 9a98d4d

2 files changed

Lines changed: 62 additions & 5 deletions

File tree

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

Lines changed: 27 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@
6161
import java.util.Optional;
6262
import java.util.Set;
6363
import java.util.concurrent.atomic.AtomicBoolean;
64+
import java.util.concurrent.atomic.AtomicInteger;
6465
import java.util.concurrent.atomic.AtomicReference;
6566
import java.util.function.BiFunction;
6667
import org.slf4j.Logger;
@@ -523,11 +524,32 @@ private Flowable<Event> runOneStep(Context spanContext, InvocationContext contex
523524
*/
524525
@Override
525526
public Flowable<Event> run(InvocationContext invocationContext) {
526-
return run(Context.current(), invocationContext, 0);
527+
Context spanContext = Context.current();
528+
return Flowable.defer(
529+
() -> {
530+
AtomicInteger stepsCompleted = new AtomicInteger(0);
531+
AtomicBoolean continueFlow = new AtomicBoolean(false);
532+
// Resubscribe without growing the stack when steps complete synchronously.
533+
return Flowable.defer(
534+
() ->
535+
runStep(
536+
spanContext,
537+
invocationContext,
538+
stepsCompleted.getAndIncrement(),
539+
continueFlow))
540+
.repeatUntil(() -> !continueFlow.getAndSet(false));
541+
});
527542
}
528543

529-
private Flowable<Event> run(
530-
Context spanContext, InvocationContext invocationContext, int stepsCompleted) {
544+
/**
545+
* Runs {@link #runOneStep} once. If the flow should run another step, sets {@code continueFlow}
546+
* and waits until the Runner has persisted this step's events.
547+
*/
548+
private Flowable<Event> runStep(
549+
Context spanContext,
550+
InvocationContext invocationContext,
551+
int stepsCompleted,
552+
AtomicBoolean continueFlow) {
531553
Flowable<Event> currentStepEvents = runOneStep(spanContext, invocationContext).cache();
532554
if (stepsCompleted + 1 >= maxSteps) {
533555
logger.debug("Ending flow execution because max steps reached.");
@@ -554,10 +576,10 @@ private Flowable<Event> run(
554576
return Flowable.empty();
555577
} else {
556578
logger.debug("Continuing to next step of the flow.");
579+
continueFlow.set(true);
557580
// Wait until the Runner has persisted this step's events so the next step's
558581
// request is not built from a stale session (see PersistBarrier).
559-
return PersistBarrier.awaitPersisted(invocationContext, eventList)
560-
.andThen(run(spanContext, invocationContext, stepsCompleted + 1));
582+
return PersistBarrier.awaitPersisted(invocationContext, eventList).toFlowable();
561583
}
562584
}));
563585
}

‎core/src/test/java/com/google/adk/flows/llmflows/BaseLlmFlowTest.java‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535
import com.google.adk.events.Event;
3636
import com.google.adk.flows.llmflows.RequestProcessor.RequestProcessingResult;
3737
import com.google.adk.flows.llmflows.ResponseProcessor.ResponseProcessingResult;
38+
import com.google.adk.models.LlmCallsLimitExceededException;
3839
import com.google.adk.models.LlmRequest;
3940
import com.google.adk.models.LlmResponse;
4041
import com.google.adk.testing.TestLlm;
@@ -43,6 +44,7 @@
4344
import com.google.adk.tools.ToolContext;
4445
import com.google.common.collect.ImmutableList;
4546
import com.google.common.collect.ImmutableMap;
47+
import com.google.common.collect.ImmutableSet;
4648
import com.google.genai.types.Content;
4749
import com.google.genai.types.FinishReason;
4850
import com.google.genai.types.FunctionCall;
@@ -59,11 +61,14 @@
5961
import io.reactivex.rxjava3.core.Maybe;
6062
import io.reactivex.rxjava3.core.Single;
6163
import io.reactivex.rxjava3.schedulers.Schedulers;
64+
import io.reactivex.rxjava3.subscribers.TestSubscriber;
65+
import java.util.ArrayList;
6266
import java.util.List;
6367
import java.util.Map;
6468
import java.util.Optional;
6569
import java.util.concurrent.atomic.AtomicInteger;
6670
import java.util.function.Function;
71+
import java.util.stream.Stream;
6772
import org.junit.Test;
6873
import org.junit.runner.RunWith;
6974
import org.junit.runners.JUnit4;
@@ -191,6 +196,36 @@ public void run_withFunctionCallsAndMaxSteps_stopsAfterMaxSteps() {
191196
Content.fromParts(Part.fromFunctionResponse("my_function", testResponse)));
192197
}
193198

199+
@Test
200+
public void run_modelRespondingOnSubscribingThread_reachesMaxLlmCallsWithoutStackOverflow() {
201+
// Flowable.just answers on the subscribing thread, like the built-in models.
202+
int maxLlmCalls = 2_000;
203+
List<Long> stackDepths = new ArrayList<>();
204+
Content contentWithFunctionCall =
205+
Content.fromParts(Part.fromFunctionCall("my_function", ImmutableMap.of("arg1", "value1")));
206+
TestLlm testLlm =
207+
createTestLlm(
208+
() -> {
209+
stackDepths.add(StackWalker.getInstance().walk(Stream::count));
210+
return Flowable.just(createLlmResponse(contentWithFunctionCall));
211+
});
212+
ImmutableMap<String, Object> testResponse =
213+
ImmutableMap.<String, Object>of("response", "response for my_function");
214+
InvocationContext invocationContext =
215+
createInvocationContext(
216+
createTestAgentBuilder(testLlm)
217+
.tools(ImmutableList.of(new TestTool("my_function", testResponse)))
218+
.build(),
219+
RunConfig.builder().maxLlmCalls(maxLlmCalls).build());
220+
BaseLlmFlow baseLlmFlow = createBaseLlmFlowWithoutProcessors();
221+
222+
TestSubscriber<Event> subscriber = baseLlmFlow.run(invocationContext).test();
223+
224+
subscriber.assertError(LlmCallsLimitExceededException.class);
225+
assertThat(testLlm.getRequests()).hasSize(maxLlmCalls);
226+
assertThat(ImmutableSet.copyOf(stackDepths)).hasSize(1);
227+
}
228+
194229
@Test
195230
public void run_withLongRunningFunctionCall_returnsCorrectEventsWithLongRunningToolIds() {
196231
Content firstContent =

0 commit comments

Comments
 (0)