Skip to content

BaseLlmFlow nests each step inside the previous one and overflows the stack after a few hundred LLM calls #1564

Description

@innoprej

🔴 Required Information

Describe the Bug:

BaseLlmFlow.run subscribes to each step (one LLM call and the tool calls it asks for) from inside the previous step's completion:

return currentStepEvents.concatWith(
    currentStepEvents
        .toList()
        .flatMapPublisher(
            eventList -> {
              ...
              return PersistBarrier.awaitPersisted(invocationContext, eventList)
                  .andThen(run(spanContext, invocationContext, stepsCompleted + 1));
            }));

If nothing in a step switches threads, the next step starts on the same call stack, 23 frames deeper than the previous one, and the stack keeps growing until the invocation ends. Nothing does with the built-in models and session services I checked: Gemini blocks the subscribing thread on the HTTP future and emits the response there (Flowable.fromFuture, with and without streaming), Claude and the non-streaming paths of LangChain4j and SpringAI return Flowable.just(...) after a blocking call, and the session services append on the calling thread (InMemorySessionService returns Single.just; VertexAiSessionService and FirestoreSessionService use Single.fromCallable without a scheduler), so the PersistBarrier wait between steps completes on the same thread too. With a 1 MB thread stack, the JVM default on x86-64, an agent that keeps calling tools fails with StackOverflowError after roughly 250–300 LLM calls in one invocation, before the default maxLlmCalls of 500 can stop it. On aarch64 the default thread stack is 2 MB: on a 2 MB stack the same agent made 635–663 calls before overflowing in my runs (on x86-64, with maxLlmCalls raised), so there the default limit usually ends the run first, and the overflow shows up with a higher maxLlmCalls or on threads with smaller stacks.

The error never reaches onError either. RxJava rethrows VirtualMachineErrors (Exceptions.throwIfFatal), so when the run is subscribed on the caller's thread, the StackOverflowError is thrown out of subscribe() and the subscriber gets neither onError nor onComplete. When the run is subscribed on a scheduler thread instead (for example with .subscribeOn(Schedulers.io())), the error goes to RxJavaPlugins.onError as an UndeliverableException and the subscriber never terminates: in my run the stack overflowed at the 254th call and blockingGet() returned only when a 15-second timeout(...) fired.

Steps to Reproduce:

  1. Check out main (4092a1f).
  2. Run an LlmAgent whose model asks for the same tool on every call and emits its response on the subscribing thread (code below) with InMemoryRunner.runAsync and the default RunConfig.
  3. On an aarch64 JVM (for example Apple silicon), add -Xss1m: its default 2 MB thread stack lets the run reach the 500-call limit first.

Expected Behavior:

The run stops at the LLM-call limit with LlmCallsLimitExceededException (500 calls by default), delivered through onError, however many steps it takes. adk-python runs the steps in a loop (while True: in BaseLlmFlow.run_async), so its stack depth does not depend on the number of steps.

Observed Behavior:

java.lang.StackOverflowError after 253 to 265 LLM calls in 14 runs of the code below with a 1 MB stack. Recording the StackWalker depth inside the test model shows the stack growing by exactly 23 frames per step, and where the error hits moves with the thread stack size:

Thread stack LLM calls before StackOverflowError (code below, in my runs)
-Xss256k 42
-Xss512k 108–113
1 MB (the x86-64 default) 253–265
2 MB (-Xss2m, the aarch64 default) none; stops at maxLlmCalls (500) as expected. With maxLlmCalls raised to 2,000: 635–663

A model or session service that completes on another thread (for example Flowable.just(...).subscribeOn(Schedulers.io()), or an appendEvent that does the same) hides the problem: each step then starts on a fresh stack, and the same agent reaches the 500-call limit.

Environment Details:

  • ADK Library Version (see maven dependency): main at 4092a1f; the v1.10.1 release has the same code. BaseLlmFlow.run has chained the steps this way since 0.1.0.
  • OS: Windows 11 on x86-64 (not OS-specific; the default thread stack size depends on the CPU architecture, see above)
  • TS Version (tsc --version): N/A (Java: Microsoft OpenJDK 17.0.19)

Model Information:

  • Which model is being used: N/A (reproduced with TestLlm; any model that emits on the subscribing thread is affected. That includes Gemini through its Flowable.fromFuture path, though I did not run it against the live API.)

🟡 Optional Information

Regression:

No.

Logs:

With the default -XX:MaxJavaStackTraceDepth (1024), the printed trace is cut off near the overflow point and shows only a five-frame loop on the request path (deepest first: SubscriptionArbiter.request, SubscriptionHelper.deferredRequest, CompletableAndThenPublisher$AndThenPublisherSubscriber.request, SubscriptionHelper.deferredRequest, SingleFlatMapPublisher$SingleFlatMapPublisherObserver.request). With -XX:MaxJavaStackTraceDepth=1000000, these are the frames that repeat once per step (RxJava 3.1.12, deepest first; package prefixes io.reactivex.rxjava3.internal.operators. and io.reactivex.rxjava3. dropped):

flowable.FlowableConcatArray$ConcatArraySubscriber.onComplete:141
flowable.FlowableConcatArray.subscribeActual:41
core.Flowable.subscribe:16149
core.Flowable.subscribe:16095
mixed.CompletableAndThenPublisher$AndThenPublisherSubscriber.onComplete:86
completable.CompletableConcatIterable$ConcatInnerObserver.next:105
completable.CompletableConcatIterable.subscribeActual:47
core.Completable.subscribe:2860
mixed.CompletableAndThenPublisher.subscribeActual:46
core.Flowable.subscribe:16149
core.Flowable.subscribe:16095
single.SingleFlatMapPublisher$SingleFlatMapPublisherObserver.onSuccess:96
flowable.FlowableToListSingle$ToListSubscriber.onComplete:102
flowable.FlowableCache.replay:245
flowable.FlowableCache$CacheSubscription.request:383
flowable.FlowableToListSingle$ToListSubscriber.onSubscribe:83
flowable.FlowableCache.subscribeActual:111
core.Flowable.subscribe:16149
flowable.FlowableToListSingle.subscribeActual:56
core.Single.subscribe:4855
single.SingleFlatMapPublisher.subscribeActual:59
core.Flowable.subscribe:16149
core.Flowable.subscribe:16095

Additional Context:

  • I have a fix that drives the steps with repeatUntil, as LoopAgent does on its resumable path, instead of subscribing to the next step from the previous one. RxJava's repeatUntil resubscribes from a trampoline loop when the source completes synchronously, so the stack depth stays the same at every step: with the fix, the code below stops after 500 calls with LlmCallsLimitExceededException even on a 256 KB stack, and a unit test runs 2,000 steps. I'll link the PR here.
  • This is not fix: avoid StackOverflowError in PersistBarrier.awaitPersisted for large steps #1336, which fixed an overflow inside PersistBarrier.awaitPersisted for a single step with many events. I first saw this overflow while working on AgentTool runs the wrapped agent with the default RunConfig instead of the caller's #1562.
  • Workaround until then: a larger thread stack (-Xss2m reached the 500-call limit in my runs) or a lower maxLlmCalls.

Minimal Reproduction Code:

// The model asks for echo_tool on every call and emits its response on the subscribing
// thread, as Gemini (Flowable.fromFuture) and Claude (Flowable.just) do.
TestLlm llm =
    createTestLlm(
        () ->
            Flowable.just(
                createFunctionCallLlmResponse("call", "echo_tool", ImmutableMap.of("arg", "x"))));
LlmAgent agent = LlmAgent.builder().name("solo").model(llm).tools(new EchoTool()).build();
InMemoryRunner runner = new InMemoryRunner(agent);
Session session = runner.sessionService().createSession(runner.appName(), "user").blockingGet();

try {
  runner
      .runAsync("user", session.id(), Content.fromParts(Part.fromText("hi")))
      .toList()
      .blockingGet();
} finally {
  System.out.println("LLM calls: " + llm.getRequests().size());
}
// main, 1 MB thread stack: throws java.lang.StackOverflowError after roughly 250-300 LLM calls.
// Expected: LlmCallsLimitExceededException after 500 calls (the default maxLlmCalls).

(createTestLlm, createFunctionCallLlmResponse and the nested EchoTool class are in TestUtils, next to TestLlm in core/src/test/java/com/google/adk/testing.)

How often has this issue occurred?:

  • Always (100%)

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions