fix:Update HITL/Tool workflows to correctly pause and resume runner operations

PiperOrigin-RevId: 840816599

Author: google-genai-bot

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

@@ -17,14 +17,18 @@
 package com.google.adk.agents;
 
 import com.google.adk.artifacts.BaseArtifactService;
+import com.google.adk.events.Event;
+import com.google.adk.flows.llmflows.ResumabilityConfig;
 import com.google.adk.memory.BaseMemoryService;
 import com.google.adk.models.LlmCallsLimitExceededException;
 import com.google.adk.plugins.PluginManager;
 import com.google.adk.sessions.BaseSessionService;
 import com.google.adk.sessions.Session;
+import com.google.common.collect.ImmutableSet;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import com.google.errorprone.annotations.InlineMe;
 import com.google.genai.types.Content;
+import com.google.genai.types.FunctionCall;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
@@ -45,6 +49,7 @@ public class InvocationContext {
   private final Session session;
   private final Optional<Content> userContent;
   private final RunConfig runConfig;
+  private final ResumabilityConfig resumabilityConfig;
   private final InvocationCostManager invocationCostManager = new InvocationCostManager();
 
   private Optional<String> branch;
@@ -64,6 +69,7 @@ private InvocationContext(Builder builder) {
     this.userContent = builder.userContent;
     this.runConfig = builder.runConfig;
     this.endInvocation = builder.endInvocation;
+    this.resumabilityConfig = builder.resumabilityConfig;
   }
 
   /**
@@ -207,6 +213,7 @@ public static InvocationContext copyOf(InvocationContext other) {
             .userContent(other.userContent)
             .runConfig(other.runConfig)
             .endInvocation(other.endInvocation)
+            .resumabilityConfig(other.resumabilityConfig)
             .build();
     newContext.activeStreamingTools.putAll(other.activeStreamingTools);
     return newContext;
@@ -248,10 +255,8 @@ public String invocationId() {
   }
 
   /**
-   * Sets the branch ID for the current invocation. A branch represents a fork in the conversation
+   * Sets the [branch] ID for the current invocation. A branch represents a fork in the conversation
    * history.
-   *
-   * @param branch the branch ID, or null to clear it
    */
   public void branch(@Nullable String branch) {
     this.branch = Optional.ofNullable(branch);
@@ -270,11 +275,7 @@ public BaseAgent agent() {
     return agent;
   }
 
-  /**
-   * Sets the agent being invoked. This is useful when delegating to a sub-agent.
-   *
-   * @param agent the agent to set
-   */
+  /** Sets the [agent] being invoked. This is useful when delegating to a sub-agent. */
   public void agent(BaseAgent agent) {
     this.agent = agent;
   }
@@ -302,11 +303,7 @@ public boolean endInvocation() {
     return endInvocation;
   }
 
-  /**
-   * Sets whether this invocation should be ended.
-   *
-   * @param endInvocation true if the invocation should end, false otherwise
-   */
+  /** Sets whether this invocation should be ended. */
   public void setEndInvocation(boolean endInvocation) {
     this.endInvocation = endInvocation;
   }
@@ -336,6 +333,28 @@ public void incrementLlmCallsCount() throws LlmCallsLimitExceededException {
     this.invocationCostManager.incrementAndEnforceLlmCallsLimit(this.runConfig);
   }
 
+  /** Returns whether the current invocation is resumable. */
+  public boolean isResumable() {
+    return resumabilityConfig.isResumable();
+  }
+
+  /** Returns whether to pause the invocation right after this [event]. */
+  public boolean shouldPauseInvocation(Event event) {
+    if (!isResumable()) {
+      return false;
+    }
+
+    var longRunningToolIds = event.longRunningToolIds().orElse(ImmutableSet.of());
+    if (longRunningToolIds.isEmpty()) {
+      return false;
+    }
+
+    return event.functionCalls().stream()
+        .map(FunctionCall::id)
+        .flatMap(Optional::stream)
+        .anyMatch(functionCallId -> longRunningToolIds.contains(functionCallId));
+  }
+
   private static class InvocationCostManager {
     private int numberOfLlmCalls = 0;
 
@@ -366,6 +385,7 @@ public static class Builder {
     private Optional<Content> userContent = Optional.empty();
     private RunConfig runConfig = RunConfig.builder().build();
     private boolean endInvocation = false;
+    private ResumabilityConfig resumabilityConfig = new ResumabilityConfig();
 
     /**
      * Sets the session service for managing session state.
@@ -553,6 +573,18 @@ public Builder endInvocation(boolean endInvocation) {
       return this;
     }
 
+    /**
+     * Sets the resumability configuration for the current agent run.
+     *
+     * @param resumabilityConfig the resumability configuration.
+     * @return this builder instance for chaining.
+     */
+    @CanIgnoreReturnValue
+    public Builder resumabilityConfig(ResumabilityConfig resumabilityConfig) {
+      this.resumabilityConfig = resumabilityConfig;
+      return this;
+    }
+
     /**
      * Builds the {@link InvocationContext} instance.
      *
@@ -584,7 +616,8 @@ public boolean equals(Object o) {
         && Objects.equals(agent, that.agent)
         && Objects.equals(session, that.session)
         && Objects.equals(userContent, that.userContent)
-        && Objects.equals(runConfig, that.runConfig);
+        && Objects.equals(runConfig, that.runConfig)
+        && Objects.equals(resumabilityConfig, that.resumabilityConfig);
   }
 
   @Override
@@ -602,6 +635,7 @@ public int hashCode() {
         session,
         userContent,
         runConfig,
-        endInvocation);
+        endInvocation,
+        resumabilityConfig);
   }
 }

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

@@ -604,7 +604,16 @@ private void maybeSaveOutputToState(Event event) {
 
   @Override
   protected Flowable<Event> runAsyncImpl(InvocationContext invocationContext) {
-    return llmFlow.run(invocationContext).doOnNext(this::maybeSaveOutputToState);
+    return llmFlow
+        .run(invocationContext)
+        .concatMap(
+            event -> {
+              this.maybeSaveOutputToState(event);
+              if (invocationContext.shouldPauseInvocation(event)) {
+                return Flowable.just(event).concatWith(Flowable.empty());
+              }
+              return Flowable.just(event);
+            });
   }
 
   @Override

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

@@ -156,36 +156,9 @@ protected Flowable<Event> postprocess(
     }
 
     return currentLlmResponse.flatMapPublisher(
-        updatedResponse -> {
-          Flowable<Event> processorEvents = Flowable.fromIterable(Iterables.concat(eventIterables));
-
-          if (updatedResponse.content().isEmpty()
-              && updatedResponse.errorCode().isEmpty()
-              && !updatedResponse.interrupted().orElse(false)
-              && !updatedResponse.turnComplete().orElse(false)) {
-            return processorEvents;
-          }
-
-          Event modelResponseEvent =
-              buildModelResponseEvent(baseEventForLlmResponse, llmRequest, updatedResponse);
-
-          Flowable<Event> modelEventStream = Flowable.just(modelResponseEvent);
-
-          if (modelResponseEvent.functionCalls().isEmpty()) {
-            return processorEvents.concatWith(modelEventStream);
-          }
-
-          Maybe<Event> maybeFunctionCallEvent;
-          if (context.runConfig().streamingMode() == StreamingMode.BIDI) {
-            maybeFunctionCallEvent =
-                Functions.handleFunctionCallsLive(context, modelResponseEvent, llmRequest.tools());
-          } else {
-            maybeFunctionCallEvent =
-                Functions.handleFunctionCalls(context, modelResponseEvent, llmRequest.tools());
-          }
-
-          return processorEvents.concatWith(modelEventStream).concatWith(maybeFunctionCallEvent);
-        });
+        updatedResponse ->
+            buildPostprocessingEvents(
+                updatedResponse, eventIterables, context, baseEventForLlmResponse, llmRequest));
   }
 
   /**
@@ -623,6 +596,45 @@ public void onError(Throwable e) {
    *
    * @return A fully constructed {@link Event} representing the LLM response.
    */
+  private Flowable<Event> buildPostprocessingEvents(
+      LlmResponse updatedResponse,
+      List<Iterable<Event>> eventIterables,
+      InvocationContext context,
+      Event baseEventForLlmResponse,
+      LlmRequest llmRequest) {
+    Flowable<Event> processorEvents = Flowable.fromIterable(Iterables.concat(eventIterables));
+    if (updatedResponse.content().isEmpty()
+        && updatedResponse.errorCode().isEmpty()
+        && !updatedResponse.interrupted().orElse(false)
+        && !updatedResponse.turnComplete().orElse(false)) {
+      return processorEvents;
+    }
+
+    Event modelResponseEvent =
+        buildModelResponseEvent(baseEventForLlmResponse, llmRequest, updatedResponse);
+    if (modelResponseEvent.functionCalls().isEmpty()) {
+      return processorEvents.concatWith(Flowable.just(modelResponseEvent));
+    }
+
+    Maybe<Event> maybeFunctionResponseEvent =
+        context.runConfig().streamingMode() == StreamingMode.BIDI
+            ? Functions.handleFunctionCallsLive(context, modelResponseEvent, llmRequest.tools())
+            : Functions.handleFunctionCalls(context, modelResponseEvent, llmRequest.tools());
+
+    Flowable<Event> functionEvents =
+        maybeFunctionResponseEvent.flatMapPublisher(
+            functionResponseEvent -> {
+              Optional<Event> toolConfirmationEvent =
+                  Functions.generateRequestConfirmationEvent(
+                      context, modelResponseEvent, functionResponseEvent);
+              return toolConfirmationEvent.isPresent()
+                  ? Flowable.just(toolConfirmationEvent.get(), functionResponseEvent)
+                  : Flowable.just(functionResponseEvent);
+            });
+
+    return processorEvents.concatWith(Flowable.just(modelResponseEvent)).concatWith(functionEvents);
+  }
+
   private Event buildModelResponseEvent(
       Event baseEventForLlmResponse, LlmRequest llmRequest, LlmResponse llmResponse) {
     Event.Builder eventBuilder =
@@ -641,10 +653,13 @@ private Event buildModelResponseEvent(
 
     Event event = eventBuilder.build();
 
+    logger.info("event: {} functionCalls: {}", event, event.functionCalls());
+
     if (!event.functionCalls().isEmpty()) {
       Functions.populateClientFunctionCallId(event);
       Set<String> longRunningToolIds =
           Functions.getLongRunningFunctionCalls(event.functionCalls(), llmRequest.tools());
+      logger.info("longRunningToolIds: {}", longRunningToolIds);
       if (!longRunningToolIds.isEmpty()) {
         event.setLongRunningToolIds(Optional.of(longRunningToolIds));
       }

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

@@ -0,0 +1,29 @@
+/*
+ * Copyright 2025 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.adk.flows.llmflows;
+
+/**
+ * An app contains Resumability configuration for the agents.
+ *
+ * @param isResumable Whether the app is resumable.
+ */
+public record ResumabilityConfig(boolean isResumable) {
+
+  /** Creates a new {@code ResumabilityConfig} with resumability disabled. */
+  public ResumabilityConfig() {
+    this(false);
+  }
+}