From 31ac2dfb920aaf3989c356186ecb0b14f503caa0 Mon Sep 17 00:00:00 2001 From: Jaehee Shin Date: Sat, 26 Sep 2026 23:11:01 +0900 Subject: [PATCH] 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. --- .../adk/flows/llmflows/BaseLlmFlow.java | 32 ++++++++++++++--- .../adk/flows/llmflows/BaseLlmFlowTest.java | 35 +++++++++++++++++++ 2 files changed, 62 insertions(+), 5 deletions(-) diff --git a/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java b/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java index b576a299e..49426a330 100644 --- a/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java +++ b/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java @@ -60,6 +60,7 @@ import java.util.Optional; import java.util.Set; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; import java.util.function.BiFunction; import org.slf4j.Logger; @@ -515,11 +516,32 @@ private Flowable runOneStep(Context spanContext, InvocationContext contex */ @Override public Flowable run(InvocationContext invocationContext) { - return run(Context.current(), invocationContext, 0); + Context spanContext = Context.current(); + return Flowable.defer( + () -> { + AtomicInteger stepsCompleted = new AtomicInteger(0); + AtomicBoolean continueFlow = new AtomicBoolean(false); + // Resubscribe without growing the stack when steps complete synchronously. + return Flowable.defer( + () -> + runStep( + spanContext, + invocationContext, + stepsCompleted.getAndIncrement(), + continueFlow)) + .repeatUntil(() -> !continueFlow.getAndSet(false)); + }); } - private Flowable run( - Context spanContext, InvocationContext invocationContext, int stepsCompleted) { + /** + * Runs {@link #runOneStep} once. If the flow should run another step, sets {@code continueFlow} + * and waits until the Runner has persisted this step's events. + */ + private Flowable runStep( + Context spanContext, + InvocationContext invocationContext, + int stepsCompleted, + AtomicBoolean continueFlow) { Flowable currentStepEvents = runOneStep(spanContext, invocationContext).cache(); if (stepsCompleted + 1 >= maxSteps) { logger.debug("Ending flow execution because max steps reached."); @@ -546,10 +568,10 @@ private Flowable run( return Flowable.empty(); } else { logger.debug("Continuing to next step of the flow."); + continueFlow.set(true); // Wait until the Runner has persisted this step's events so the next step's // request is not built from a stale session (see PersistBarrier). - return PersistBarrier.awaitPersisted(invocationContext, eventList) - .andThen(run(spanContext, invocationContext, stepsCompleted + 1)); + return PersistBarrier.awaitPersisted(invocationContext, eventList).toFlowable(); } })); } diff --git a/core/src/test/java/com/google/adk/flows/llmflows/BaseLlmFlowTest.java b/core/src/test/java/com/google/adk/flows/llmflows/BaseLlmFlowTest.java index 7da99cef0..b03cf4695 100644 --- a/core/src/test/java/com/google/adk/flows/llmflows/BaseLlmFlowTest.java +++ b/core/src/test/java/com/google/adk/flows/llmflows/BaseLlmFlowTest.java @@ -35,6 +35,7 @@ import com.google.adk.events.Event; import com.google.adk.flows.llmflows.RequestProcessor.RequestProcessingResult; import com.google.adk.flows.llmflows.ResponseProcessor.ResponseProcessingResult; +import com.google.adk.models.LlmCallsLimitExceededException; import com.google.adk.models.LlmRequest; import com.google.adk.models.LlmResponse; import com.google.adk.testing.TestLlm; @@ -43,6 +44,7 @@ import com.google.adk.tools.ToolContext; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; import com.google.genai.types.Content; import com.google.genai.types.FinishReason; import com.google.genai.types.FunctionCall; @@ -59,11 +61,14 @@ import io.reactivex.rxjava3.core.Maybe; import io.reactivex.rxjava3.core.Single; import io.reactivex.rxjava3.schedulers.Schedulers; +import io.reactivex.rxjava3.subscribers.TestSubscriber; +import java.util.ArrayList; import java.util.List; import java.util.Map; import java.util.Optional; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Function; +import java.util.stream.Stream; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -191,6 +196,36 @@ public void run_withFunctionCallsAndMaxSteps_stopsAfterMaxSteps() { Content.fromParts(Part.fromFunctionResponse("my_function", testResponse))); } + @Test + public void run_modelRespondingOnSubscribingThread_reachesMaxLlmCallsWithoutStackOverflow() { + // Flowable.just answers on the subscribing thread, like the built-in models. + int maxLlmCalls = 2_000; + List stackDepths = new ArrayList<>(); + Content contentWithFunctionCall = + Content.fromParts(Part.fromFunctionCall("my_function", ImmutableMap.of("arg1", "value1"))); + TestLlm testLlm = + createTestLlm( + () -> { + stackDepths.add(StackWalker.getInstance().walk(Stream::count)); + return Flowable.just(createLlmResponse(contentWithFunctionCall)); + }); + ImmutableMap testResponse = + ImmutableMap.of("response", "response for my_function"); + InvocationContext invocationContext = + createInvocationContext( + createTestAgentBuilder(testLlm) + .tools(ImmutableList.of(new TestTool("my_function", testResponse))) + .build(), + RunConfig.builder().maxLlmCalls(maxLlmCalls).build()); + BaseLlmFlow baseLlmFlow = createBaseLlmFlowWithoutProcessors(); + + TestSubscriber subscriber = baseLlmFlow.run(invocationContext).test(); + + subscriber.assertError(LlmCallsLimitExceededException.class); + assertThat(testLlm.getRequests()).hasSize(maxLlmCalls); + assertThat(ImmutableSet.copyOf(stackDepths)).hasSize(1); + } + @Test public void run_withLongRunningFunctionCall_returnsCorrectEventsWithLongRunningToolIds() { Content firstContent =