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 =