feat: introduce an experimental parameter to limit number of steps LlmAgent can take

PiperOrigin-RevId: 778386550

Author: google-genai-bot

core/src/main/java/com/google/adk/agents/LlmAgent.java

@@ -87,6 +87,7 @@ public enum IncludeContents {
   private final IncludeContents includeContents;
 
   private final boolean planning;
+  private final Optional<Integer> maxSteps;
   private final boolean disallowTransferToParent;
   private final boolean disallowTransferToPeers;
   private final Optional<List<BeforeModelCallback>> beforeModelCallback;
@@ -118,6 +119,7 @@ protected LlmAgent(Builder builder) {
     this.includeContents =
         builder.includeContents != null ? builder.includeContents : IncludeContents.DEFAULT;
     this.planning = builder.planning != null && builder.planning;
+    this.maxSteps = Optional.ofNullable(builder.maxSteps);
     this.disallowTransferToParent = builder.disallowTransferToParent;
     this.disallowTransferToPeers = builder.disallowTransferToPeers;
     this.beforeModelCallback = Optional.ofNullable(builder.beforeModelCallback);
@@ -156,6 +158,7 @@ public static class Builder {
     private BaseExampleProvider exampleProvider;
     private IncludeContents includeContents;
     private Boolean planning;
+    private Integer maxSteps;
     private Boolean disallowTransferToParent;
     private Boolean disallowTransferToPeers;
     private ImmutableList<BeforeModelCallback> beforeModelCallback;
@@ -290,6 +293,12 @@ public Builder planning(boolean planning) {
       return this;
     }
 
+    @CanIgnoreReturnValue
+    public Builder maxSteps(int maxSteps) {
+      this.maxSteps = maxSteps;
+      return this;
+    }
+
     @CanIgnoreReturnValue
     public Builder disallowTransferToParent(boolean disallowTransferToParent) {
       this.disallowTransferToParent = disallowTransferToParent;
@@ -588,9 +597,9 @@ public LlmAgent build() {
 
   protected BaseLlmFlow determineLlmFlow() {
     if (disallowTransferToParent() && disallowTransferToPeers() && subAgents().isEmpty()) {
-      return new SingleFlow();
+      return new SingleFlow(maxSteps);
     } else {
-      return new AutoFlow();
+      return new AutoFlow(maxSteps);
     }
   }
 

core/src/main/java/com/google/adk/flows/llmflows/AutoFlow.java

@@ -17,6 +17,7 @@
 package com.google.adk.flows.llmflows;
 
 import com.google.common.collect.ImmutableList;
+import java.util.Optional;
 
 /** LLM flow with automatic agent transfer support. */
 public class AutoFlow extends SingleFlow {
@@ -32,6 +33,10 @@ public class AutoFlow extends SingleFlow {
   private static final ImmutableList<ResponseProcessor> RESPONSE_PROCESSORS = ImmutableList.of();
 
   public AutoFlow() {
-    super(REQUEST_PROCESSORS, RESPONSE_PROCESSORS);
+    this(/* maxSteps= */ Optional.empty());
+  }
+
+  public AutoFlow(Optional<Integer> maxSteps) {
+    super(REQUEST_PROCESSORS, RESPONSE_PROCESSORS, maxSteps);
   }
 }

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

@@ -65,10 +65,23 @@ public abstract class BaseLlmFlow implements BaseFlow {
   protected final List<RequestProcessor> requestProcessors;
   protected final List<ResponseProcessor> responseProcessors;
 
+  // Warning: This is local, in-process state that won't be preserved if the runtime is restarted.
+  // "Max steps" is experimental and may evolve in the future (e.g., to support persistence).
+  protected int stepsCompleted = 0;
+  protected final int maxSteps;
+
   public BaseLlmFlow(
       List<RequestProcessor> requestProcessors, List<ResponseProcessor> responseProcessors) {
+    this(requestProcessors, responseProcessors, /* maxSteps= */ Optional.empty());
+  }
+
+  public BaseLlmFlow(
+      List<RequestProcessor> requestProcessors,
+      List<ResponseProcessor> responseProcessors,
+      Optional<Integer> maxSteps) {
     this.requestProcessors = requestProcessors;
     this.responseProcessors = responseProcessors;
+    this.maxSteps = maxSteps.orElse(Integer.MAX_VALUE);
   }
 
   /**
@@ -407,6 +420,11 @@ private Flowable<Event> runOneStep(InvocationContext context) {
   @Override
   public Flowable<Event> run(InvocationContext invocationContext) {
     Flowable<Event> currentStepEvents = runOneStep(invocationContext).cache();
+    if (++stepsCompleted >= maxSteps) {
+      logger.debug("Ending flow execution because max steps reached.");
+      return currentStepEvents;
+    }
+
     return currentStepEvents.concatWith(
         currentStepEvents
             .toList()

core/src/main/java/com/google/adk/flows/llmflows/SingleFlow.java

@@ -18,6 +18,7 @@
 
 import com.google.common.collect.ImmutableList;
 import java.util.List;
+import java.util.Optional;
 
 /** Basic LLM flow with fixed request processors and no response post-processing. */
 public class SingleFlow extends BaseLlmFlow {
@@ -29,11 +30,17 @@ public class SingleFlow extends BaseLlmFlow {
   protected static final ImmutableList<ResponseProcessor> RESPONSE_PROCESSORS = ImmutableList.of();
 
   public SingleFlow() {
-    super(REQUEST_PROCESSORS, RESPONSE_PROCESSORS);
+    this(/* maxSteps= */ Optional.empty());
+  }
+
+  public SingleFlow(Optional<Integer> maxSteps) {
+    this(REQUEST_PROCESSORS, RESPONSE_PROCESSORS, maxSteps);
   }
 
   protected SingleFlow(
-      List<RequestProcessor> requestProcessors, List<ResponseProcessor> responseProcessors) {
-    super(requestProcessors, responseProcessors);
+      List<RequestProcessor> requestProcessors,
+      List<ResponseProcessor> responseProcessors,
+      Optional<Integer> maxSteps) {
+    super(requestProcessors, responseProcessors, maxSteps);
   }
 }

core/src/test/java/com/google/adk/agents/LlmAgentTest.java

@@ -16,6 +16,7 @@
 
 package com.google.adk.agents;
 
+import static com.google.adk.testing.TestUtils.assertEqualIgnoringFunctionIds;
 import static com.google.adk.testing.TestUtils.createInvocationContext;
 import static com.google.adk.testing.TestUtils.createLlmResponse;
 import static com.google.adk.testing.TestUtils.createTestAgent;
@@ -28,6 +29,7 @@
 import com.google.adk.events.Event;
 import com.google.adk.models.LlmResponse;
 import com.google.adk.testing.TestLlm;
+import com.google.adk.testing.TestUtils.EchoTool;
 import com.google.adk.tools.BaseTool;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
@@ -108,6 +110,33 @@ public void testRun_withoutOutputKey_doesNotSaveState() {
     assertThat(events.get(0).actions().stateDelta()).isEmpty();
   }
 
+  @Test
+  public void run_withToolsAndMaxSteps_stopsAfterMaxSteps() {
+    ImmutableMap<String, Object> echoArgs = ImmutableMap.of("arg", "value");
+    Content contentWithFunctionCall =
+        Content.fromParts(Part.fromText("text"), Part.fromFunctionCall("echo_tool", echoArgs));
+    Content unreachableContent = Content.fromParts(Part.fromText("This should never be returned."));
+    TestLlm testLlm =
+        createTestLlm(
+            createLlmResponse(contentWithFunctionCall),
+            createLlmResponse(contentWithFunctionCall),
+            createLlmResponse(unreachableContent));
+    LlmAgent agent = createTestAgentBuilder(testLlm).tools(new EchoTool()).maxSteps(2).build();
+    InvocationContext invocationContext = createInvocationContext(agent);
+
+    List<Event> events = agent.runAsync(invocationContext).toList().blockingGet();
+
+    Content expectedFunctionResponseContent =
+        Content.fromParts(
+            Part.fromFunctionResponse(
+                "echo_tool", ImmutableMap.<String, Object>of("result", echoArgs)));
+    assertThat(events).hasSize(4);
+    assertEqualIgnoringFunctionIds(events.get(0).content().get(), contentWithFunctionCall);
+    assertEqualIgnoringFunctionIds(events.get(1).content().get(), expectedFunctionResponseContent);
+    assertEqualIgnoringFunctionIds(events.get(2).content().get(), contentWithFunctionCall);
+    assertEqualIgnoringFunctionIds(events.get(3).content().get(), expectedFunctionResponseContent);
+  }
+
   @Test
   public void build_withOutputSchemaAndTools_throwsIllegalArgumentException() {
     BaseTool tool =

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

@@ -16,6 +16,7 @@
 
 package com.google.adk.flows.llmflows;
 
+import static com.google.adk.testing.TestUtils.assertEqualIgnoringFunctionIds;
 import static com.google.adk.testing.TestUtils.createInvocationContext;
 import static com.google.adk.testing.TestUtils.createLlmResponse;
 import static com.google.adk.testing.TestUtils.createTestAgent;
@@ -36,9 +37,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.genai.types.Content;
-import com.google.genai.types.FunctionCall;
 import com.google.genai.types.FunctionDeclaration;
-import com.google.genai.types.FunctionResponse;
 import com.google.genai.types.Part;
 import io.reactivex.rxjava3.core.Flowable;
 import io.reactivex.rxjava3.core.Single;
@@ -91,21 +90,51 @@ public void run_withFunctionCall_returnsCorrectEvents() {
     List<Event> events = baseLlmFlow.run(invocationContext).toList().blockingGet();
 
     assertThat(events).hasSize(3);
-    assertContentIgnoringFunctionId(events.get(0).content().get(), firstContent);
-    assertContentIgnoringFunctionId(
+    assertEqualIgnoringFunctionIds(events.get(0).content().get(), firstContent);
+    assertEqualIgnoringFunctionIds(
         events.get(1).content().get(),
-        Content.fromParts(
-            Part.builder()
-                .functionResponse(
-                    FunctionResponse.builder()
-                        .id("")
-                        .name("my_function")
-                        .response(testResponse)
-                        .build())
-                .build()));
+        Content.fromParts(Part.fromFunctionResponse("my_function", testResponse)));
     assertThat(events.get(2).content()).hasValue(secondContent);
   }
 
+  @Test
+  public void run_withFunctionCallsAndMaxSteps_stopsAfterMaxSteps() {
+    Content contentWithFunctionCall =
+        Content.fromParts(
+            Part.fromText("LLM response with function call"),
+            Part.fromFunctionCall("my_function", ImmutableMap.of("arg1", "value1")));
+    Content unreachableContent = Content.fromParts(Part.fromText("This should never be returned."));
+    TestLlm testLlm =
+        createTestLlm(
+            Flowable.just(createLlmResponse(contentWithFunctionCall)),
+            Flowable.just(createLlmResponse(contentWithFunctionCall)),
+            Flowable.just(createLlmResponse(unreachableContent)));
+    ImmutableMap<String, Object> testResponse =
+        ImmutableMap.<String, Object>of("response", "response for my_function");
+    InvocationContext invocationContext =
+        createInvocationContext(
+            createTestAgentBuilder(testLlm)
+                .tools(ImmutableList.of(new TestTool("my_function", testResponse)))
+                .build());
+    BaseLlmFlow baseLlmFlow =
+        createBaseLlmFlow(
+            /* requestProcessors= */ ImmutableList.of(),
+            /* responseProcessors= */ ImmutableList.of(),
+            /* maxSteps= */ Optional.of(2));
+
+    List<Event> events = baseLlmFlow.run(invocationContext).toList().blockingGet();
+
+    assertThat(events).hasSize(4);
+    assertEqualIgnoringFunctionIds(events.get(0).content().get(), contentWithFunctionCall);
+    assertEqualIgnoringFunctionIds(
+        events.get(1).content().get(),
+        Content.fromParts(Part.fromFunctionResponse("my_function", testResponse)));
+    assertEqualIgnoringFunctionIds(events.get(2).content().get(), contentWithFunctionCall);
+    assertEqualIgnoringFunctionIds(
+        events.get(3).content().get(),
+        Content.fromParts(Part.fromFunctionResponse("my_function", testResponse)));
+  }
+
   @Test
   public void run_withRequestProcessor_doesNotModifyRequest() {
     Content content = Content.fromParts(Part.fromText("LLM response"));
@@ -206,7 +235,15 @@ private static BaseLlmFlow createBaseLlmFlowWithoutProcessors() {
 
   private static BaseLlmFlow createBaseLlmFlow(
       List<RequestProcessor> requestProcessors, List<ResponseProcessor> responseProcessors) {
-    return new BaseLlmFlow(requestProcessors, responseProcessors) {};
+    return createBaseLlmFlow(
+        requestProcessors, responseProcessors, /* maxSteps= */ Optional.empty());
+  }
+
+  private static BaseLlmFlow createBaseLlmFlow(
+      List<RequestProcessor> requestProcessors,
+      List<ResponseProcessor> responseProcessors,
+      Optional<Integer> maxSteps) {
+    return new BaseLlmFlow(requestProcessors, responseProcessors, maxSteps) {};
   }
 
   private static RequestProcessor createRequestProcessor() {
@@ -256,41 +293,4 @@ public Single<Map<String, Object>> runAsync(Map<String, Object> args, ToolContex
       return Single.just(response);
     }
   }
-
-  private static void assertContentIgnoringFunctionId(
-      Content actualContent, Content expectedContent) {
-
-    assertThat(actualContent.role()).isEqualTo(expectedContent.role());
-
-    Optional<List<Part>> actualPartsOpt = actualContent.parts();
-    Optional<List<Part>> expectedPartsOpt = expectedContent.parts();
-    assertThat(actualPartsOpt.isPresent()).isEqualTo(expectedPartsOpt.isPresent());
-
-    if (expectedPartsOpt.isPresent()) {
-      List<Part> actualParts = actualPartsOpt.get();
-      List<Part> expectedParts = expectedPartsOpt.get();
-      assertThat(actualParts).hasSize(expectedParts.size());
-
-      for (int i = 0; i < expectedParts.size(); i++) {
-        Part actualPart = actualParts.get(i);
-        Part expectedPart = expectedParts.get(i);
-
-        if (expectedPart.functionCall().isPresent()) {
-          assertThat(actualPart.functionCall()).isPresent();
-          FunctionCall actualFc = actualPart.functionCall().get();
-          FunctionCall expectedFc = expectedPart.functionCall().get();
-          assertThat(actualFc.name()).isEqualTo(expectedFc.name());
-          assertThat(actualFc.args()).isEqualTo(expectedFc.args());
-        } else if (expectedPart.functionResponse().isPresent()) {
-          assertThat(actualPart.functionResponse()).isPresent();
-          FunctionResponse actualFr = actualPart.functionResponse().get();
-          FunctionResponse expectedFr = expectedPart.functionResponse().get();
-          assertThat(actualFr.name()).isEqualTo(expectedFr.name());
-          assertThat(actualFr.response()).isEqualTo(expectedFr.response());
-        } else {
-          assertThat(actualPart).isEqualTo(expectedPart);
-        }
-      }
-    }
-  }
 }

core/src/test/java/com/google/adk/testing/TestUtils.java

@@ -17,6 +17,7 @@
 package com.google.adk.testing;
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
+import static com.google.common.truth.Truth.assertThat;
 import static java.util.stream.Collectors.joining;
 
 import com.google.adk.agents.BaseAgent;
@@ -142,6 +143,40 @@ public static ImmutableList<Object> simplifyEvents(List<Event> events) {
         .collect(toImmutableList());
   }
 
+  public static void assertEqualIgnoringFunctionIds(
+      Content actualContent, Content expectedContent) {
+    assertThat(overwriteFunctionIdsInContent(actualContent))
+        .isEqualTo(overwriteFunctionIdsInContent(expectedContent));
+  }
+
+  private static Content overwriteFunctionIdsInContent(Content content) {
+    if (content.parts().isEmpty()) {
+      return content;
+    }
+    return content.toBuilder()
+        .parts(
+            content.parts().get().stream()
+                .map(TestUtils::overwriteFunctionIdsInPart)
+                .collect(toImmutableList()))
+        .build();
+  }
+
+  private static Part overwriteFunctionIdsInPart(Part part) {
+    if (part.functionCall().isPresent()) {
+      return part.toBuilder()
+          .functionCall(
+              part.functionCall().get().toBuilder().id("<overwritten by TestUtils>").build())
+          .build();
+    }
+    if (part.functionResponse().isPresent()) {
+      return part.toBuilder()
+          .functionResponse(
+              part.functionResponse().get().toBuilder().id("<overwritten by TestUtils>").build())
+          .build();
+    }
+    return part;
+  }
+
   public static TestBaseAgent createRootAgent(BaseAgent... subAgents) {
     return createRootAgent(Arrays.asList(subAgents));
   }