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

PiperOrigin-RevId: 833856048

Author: google-genai-bot

core/src/main/java/com/google/adk/agents/InvocationContext.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/agents/LlmAgent.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 =
@@ -640,10 +652,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/BaseLlmFlow.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);
+  }
+}

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

@@ -26,6 +26,7 @@
 import com.google.adk.artifacts.BaseArtifactService;
 import com.google.adk.events.Event;
 import com.google.adk.events.EventActions;
+import com.google.adk.flows.llmflows.ResumabilityConfig;
 import com.google.adk.memory.BaseMemoryService;
 import com.google.adk.plugins.BasePlugin;
 import com.google.adk.plugins.PluginManager;
@@ -62,8 +63,9 @@ public class Runner {
   private final String appName;
   private final BaseArtifactService artifactService;
   private final BaseSessionService sessionService;
-  private final @Nullable BaseMemoryService memoryService;
+  @Nullable private final BaseMemoryService memoryService;
   private final PluginManager pluginManager;
+  private final ResumabilityConfig resumabilityConfig;
 
   /** Creates a new {@code Runner}. */
   public Runner(
@@ -72,7 +74,14 @@ public Runner(
       BaseArtifactService artifactService,
       BaseSessionService sessionService,
       @Nullable BaseMemoryService memoryService) {
-    this(agent, appName, artifactService, sessionService, memoryService, ImmutableList.of());
+    this(
+        agent,
+        appName,
+        artifactService,
+        sessionService,
+        memoryService,
+        ImmutableList.of(),
+        new ResumabilityConfig());
   }
 
   /** Creates a new {@code Runner} with a list of plugins. */
@@ -83,12 +92,32 @@ public Runner(
       BaseSessionService sessionService,
       @Nullable BaseMemoryService memoryService,
       List<BasePlugin> plugins) {
+    this(
+        agent,
+        appName,
+        artifactService,
+        sessionService,
+        memoryService,
+        plugins,
+        new ResumabilityConfig());
+  }
+
+  /** Creates a new {@code Runner} with a list of plugins and resumability config. */
+  public Runner(
+      BaseAgent agent,
+      String appName,
+      BaseArtifactService artifactService,
+      BaseSessionService sessionService,
+      @Nullable BaseMemoryService memoryService,
+      List<BasePlugin> plugins,
+      ResumabilityConfig resumabilityConfig) {
     this.agent = agent;
     this.appName = appName;
     this.artifactService = artifactService;
     this.sessionService = sessionService;
     this.memoryService = memoryService;
     this.pluginManager = new PluginManager(plugins);
+    this.resumabilityConfig = resumabilityConfig;
   }
 
   /**
@@ -123,7 +152,8 @@ public BaseSessionService sessionService() {
     return this.sessionService;
   }
 
-  public @Nullable BaseMemoryService memoryService() {
+  @Nullable
+  public BaseMemoryService memoryService() {
     return this.memoryService;
   }
 
@@ -305,7 +335,7 @@ public Flowable<Event> runAsync(
                                                   newInvocationContextWithId(
                                                       updatedSession,
                                                       event.content(),
-                                                      Optional.empty(),
+                                                      /* liveRequestQueue= */ Optional.empty(),
                                                       runConfig,
                                                       invocationId);
                                               contextWithUpdatedSession.agent(
@@ -430,20 +460,19 @@ private InvocationContext newInvocationContext(
       Optional<LiveRequestQueue> liveRequestQueue,
       RunConfig runConfig) {
     BaseAgent rootAgent = this.agent;
-    InvocationContext invocationContext =
-        new InvocationContext(
-            this.sessionService,
-            this.artifactService,
-            this.memoryService,
-            this.pluginManager,
-            liveRequestQueue,
-            /* branch= */ Optional.empty(),
-            InvocationContext.newInvocationContextId(),
-            rootAgent,
-            session,
-            newMessage,
-            runConfig,
-            /* endInvocation= */ false);
+    var invocationContextBuilder =
+        InvocationContext.builder()
+            .sessionService(this.sessionService)
+            .artifactService(this.artifactService)
+            .memoryService(this.memoryService)
+            .pluginManager(this.pluginManager)
+            .agent(rootAgent)
+            .session(session)
+            .userContent(newMessage)
+            .runConfig(runConfig)
+            .resumabilityConfig(this.resumabilityConfig);
+    liveRequestQueue.ifPresent(invocationContextBuilder::liveRequestQueue);
+    var invocationContext = invocationContextBuilder.build();
     invocationContext.agent(this.findAgentToRun(session, rootAgent));
     return invocationContext;
   }
@@ -460,20 +489,20 @@ private InvocationContext newInvocationContextWithId(
       RunConfig runConfig,
       String invocationId) {
     BaseAgent rootAgent = this.agent;
-    InvocationContext invocationContext =
-        new InvocationContext(
-            this.sessionService,
-            this.artifactService,
-            this.memoryService,
-            this.pluginManager,
-            liveRequestQueue,
-            /* branch= */ Optional.empty(),
-            invocationId,
-            rootAgent,
-            session,
-            newMessage,
-            runConfig,
-            /* endInvocation= */ false);
+    var invocationContextBuilder =
+        InvocationContext.builder()
+            .sessionService(this.sessionService)
+            .artifactService(this.artifactService)
+            .memoryService(this.memoryService)
+            .pluginManager(this.pluginManager)
+            .invocationId(invocationId)
+            .agent(rootAgent)
+            .session(session)
+            .userContent(newMessage)
+            .runConfig(runConfig)
+            .resumabilityConfig(this.resumabilityConfig);
+    liveRequestQueue.ifPresent(invocationContextBuilder::liveRequestQueue);
+    var invocationContext = invocationContextBuilder.build();
     invocationContext.agent(this.findAgentToRun(session, rootAgent));
     return invocationContext;
   }

core/src/main/java/com/google/adk/runner/Runner.java

@@ -17,40 +17,65 @@
 package com.google.adk.tools;
 
 import java.lang.reflect.Method;
+import javax.annotation.Nullable;
 
 /** A function tool that returns the result asynchronously. */
 public class LongRunningFunctionTool extends FunctionTool {
 
   public static LongRunningFunctionTool create(Method func) {
-    return new LongRunningFunctionTool(func);
+    return create(func, /* requireConfirmation= */ false);
+  }
+
+  public static LongRunningFunctionTool create(Method func, boolean requireConfirmation) {
+    return new LongRunningFunctionTool(func, requireConfirmation);
   }
 
   public static LongRunningFunctionTool create(Class<?> cls, String methodName) {
+    return create(cls, methodName, /* requireConfirmation= */ false);
+  }
+
+  public static LongRunningFunctionTool create(
+      Class<?> cls, String methodName, boolean requireConfirmation) {
     for (Method method : cls.getMethods()) {
       if (method.getName().equals(methodName)) {
-        return create(method);
+        return create(method, requireConfirmation);
       }
     }
     throw new IllegalArgumentException(
         String.format("Method %s not found in class %s.", methodName, cls.getName()));
   }
 
   public static LongRunningFunctionTool create(Object instance, String methodName) {
+    return create(instance, methodName, /* requireConfirmation= */ false);
+  }
+
+  public static LongRunningFunctionTool create(
+      Object instance, String methodName, boolean requireConfirmation) {
     Class<?> cls = instance.getClass();
     for (Method method : cls.getMethods()) {
       if (method.getName().equals(methodName)) {
-        return new LongRunningFunctionTool(instance, method);
+        return create(instance, method, requireConfirmation);
       }
     }
     throw new IllegalArgumentException(
         String.format("Method %s not found in class %s.", methodName, cls.getName()));
   }
 
-  private LongRunningFunctionTool(Method func) {
-    super(null, func, /* isLongRunning= */ true, /* requireConfirmation= */ false);
+  public static LongRunningFunctionTool create(@Nullable Object instance, Method method) {
+    return create(instance, method, false);
+  }
+
+  public static LongRunningFunctionTool create(
+      @Nullable Object instance, Method method, boolean requireConfirmation) {
+    return new LongRunningFunctionTool(instance, method, requireConfirmation);
+  }
+
+  private LongRunningFunctionTool(Method func, boolean requireConfirmation) {
+    super(null, func, /* isLongRunning= */ true, requireConfirmation);
   }
 
-  private LongRunningFunctionTool(Object instance, Method func) {
-    super(instance, func, /* isLongRunning= */ true, /* requireConfirmation= */ false);
+  private LongRunningFunctionTool(
+      @Nullable Object instance, Method func, boolean requireConfirmation) {
+    super(instance, func, /* isLongRunning= */ true, requireConfirmation);
   }
 }