feat: share one run per Context.runNode run ID

PiperOrigin-RevId: 1004727068

Author: damianmomotgoogle

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

@@ -35,6 +35,7 @@
 import io.reactivex.rxjava3.core.Maybe;
 import io.reactivex.rxjava3.core.Single;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import org.jspecify.annotations.Nullable;
 
@@ -321,6 +322,18 @@ public final void routes(List<? extends Route> routes) {
     requireNodeState().selectRoutes(routes);
   }
 
+  /**
+   * Returns the user's answers to the long-running calls that this workflow node or its descendants
+   * made in an earlier turn, keyed by interrupt ID. A node that reruns on resume reads these
+   * answers. A null answer has a null value in the map.
+   *
+   * @throws IllegalStateException if this is not the context of a workflow node.
+   */
+  @Experimental
+  public final Map<String, @Nullable Object> resumeInputs() {
+    return requireNodeState().resumeInputs();
+  }
+
   /**
    * Runs {@code node} as a child of this workflow node without an input, as {@link #runNode(Node,
    * Object, RunNodeConfig)} does with the default {@link RunNodeConfig}.

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

@@ -43,6 +43,7 @@
 import com.google.adk.agents.Callbacks.OnToolErrorCallbackBase;
 import com.google.adk.agents.Callbacks.OnToolErrorCallbackSync;
 import com.google.adk.agents.ConfigAgentUtils.ConfigurationException;
+import com.google.adk.annotations.Experimental;
 import com.google.adk.codeexecutors.BaseCodeExecutor;
 import com.google.adk.events.Event;
 import com.google.adk.flows.llmflows.AutoFlow;
@@ -1082,6 +1083,16 @@ public Optional<Schema> outputSchema() {
     return outputSchema;
   }
 
+  /**
+   * Returns {@code true}, so that a workflow resuming this agent as a node runs it again and its
+   * model sees the answers to its long-running calls.
+   */
+  @Experimental
+  @Override
+  public boolean rerunOnResume() {
+    return true;
+  }
+
   public Optional<Executor> executor() {
     return executor;
   }

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

@@ -111,8 +111,6 @@ public class Runner {
   private final UuidProvider uuidProvider;
   private final ConcurrentMap<String, Completable> activeSessionCompletables =
       new MapMaker().weakValues().makeMap();
-  private static final String NODE_ROOT_RESUME_UNSUPPORTED =
-      "Resuming an invocation of a workflow node root is not supported.";
 
   /** Builder for {@link Runner}. */
   public static class Builder {
@@ -668,7 +666,8 @@ public Flowable<Event> runAsync(String userId, String sessionId, Content newMess
    * Runs the agent, resuming an existing invocation instead of starting a new one. The invocation
    * is resolved from {@code invocationId}, or from a function response carried by {@code
    * newMessage}. Agent checkpoints are rehydrated from history and an invocation whose active agent
-   * already finished resolves to a no-op.
+   * already finished resolves to a no-op, while a workflow node root rebuilds its progress from the
+   * invocation's events even when the app is not resumable.
    *
    * @param userId the user id of the session.
    * @param sessionId the session id.
@@ -679,7 +678,7 @@ public Flowable<Event> runAsync(String userId, String sessionId, Content newMess
    * @param stateDelta optional state updates to merge into the session for this run.
    * @return the events generated while resuming, or an empty stream when there is nothing to
    *     resume.
-   * @throws IllegalStateException if the app is not resumable.
+   * @throws IllegalStateException if the app is not resumable and its root is an agent.
    * @throws IllegalArgumentException if the invocation cannot be resolved.
    */
   @Experimental
@@ -691,7 +690,7 @@ public Flowable<Event> runAsync(
       RunConfig runConfig,
       @Nullable Map<String, Object> stateDelta) {
     checkState(
-        isResumable(),
+        isResumable() || node != null,
         "Resuming an invocation requires an App configured with a resumable ResumabilityConfig.");
     Flowable<Event> result =
         Flowable.defer(
@@ -782,26 +781,25 @@ protected Flowable<Event> runAsyncImpl(
     Preconditions.checkNotNull(session, "session cannot be null");
     Preconditions.checkNotNull(newMessage, "newMessage cannot be null");
     Preconditions.checkNotNull(runConfig, "runConfig cannot be null");
-    if (isResumable()) {
+    // A node root rebuilds its progress from its invocation's events, so it resumes on any app.
+    if (isResumable() || node != null) {
       return runResumableFromSession(
           session, /* providedInvocationId= */ null, newMessage, runConfig, stateDelta);
     }
     return runNewInvocation(session, newMessage, runConfig, stateDelta);
   }
 
-  /** Starts a brand-new invocation for {@code newMessage} (the default, non-resume flow). */
+  /**
+   * Starts an invocation for {@code newMessage}; a node root joins the one a plugin's answers
+   * resolve to.
+   */
   private Flowable<Event> runNewInvocation(
       Session session,
       Content newMessage,
       RunConfig runConfig,
       @Nullable Map<String, Object> stateDelta) {
     return Flowable.defer(
             () -> {
-              // Rejected before any callback or append; defer turns the throw into an error signal.
-              if (node != null) {
-                checkNoTextWithFunctionResponses(newMessage);
-                checkNoAnswerToAnEarlierInvocation(session, newMessage);
-              }
               Context capturedContext = Context.current();
               BaseAgent rootAgent = this.agent;
               String invocationId = InvocationContext.newInvocationContextId(this.uuidProvider);
@@ -825,18 +823,21 @@ private Flowable<Event> runNewInvocation(
                   .onUserMessageCallback(initialContext, newMessage)
                   .compose(Tracing.<Content>withContext(capturedContext))
                   .defaultIfEmpty(newMessage)
-                  .flatMap(
-                      content ->
-                          createUserMessageEvent(
-                              session,
-                              content,
-                              initialContext,
-                              runConfig.saveInputBlobsAsArtifacts(),
-                              stateDelta))
                   .flatMapPublisher(
-                      userEvent ->
-                          runAgentForUserEvent(initialContext, session, userEvent, rootAgent)
-                              .compose(Tracing.<Event>withContext(capturedContext)))
+                      content -> {
+                        InvocationContext context =
+                            resolveNodeRootContext(session, initialContext, content);
+                        return createUserMessageEvent(
+                                session,
+                                content,
+                                context,
+                                runConfig.saveInputBlobsAsArtifacts(),
+                                stateDelta)
+                            .flatMapPublisher(
+                                userEvent ->
+                                    runAgentForUserEvent(context, session, userEvent, rootAgent)
+                                        .compose(Tracing.<Event>withContext(capturedContext)));
+                      })
                   .doOnError(throwable -> runOnRunError(initialContext, throwable));
             })
         .doOnError(
@@ -893,7 +894,11 @@ private Flowable<Event> runAgentWithUpdatedSession(
         initialContext.toBuilder()
             .session(updatedSession)
             .agent(agentToRun)
-            .userContent(event.content().orElseGet(Content::fromParts))
+            .userContent(
+                userContentFor(
+                    updatedSession,
+                    initialContext.invocationId(),
+                    event.content().orElseGet(Content::fromParts)))
             .branch(routedAgentParentBranch(updatedSession, agentToRun, rootAgent))
             .build();
 
@@ -1006,7 +1011,7 @@ private void runOnRunError(InvocationContext context, Throwable throwable) {
   /**
    * Resumes an existing invocation when one resolves from {@code providedInvocationId} or a
    * function response in {@code rawMessage}; otherwise starts a new invocation. Requires
-   * resumability.
+   * resumability or a workflow node root.
    */
   private Flowable<Event> runResumableFromSession(
       Session session,
@@ -1043,11 +1048,6 @@ private Flowable<Event> runResumableFromSession(
             return Flowable.<Event>error(
                 new IllegalArgumentException("Session has no events to resume."));
           }
-          if (node != null) {
-            // Without replay, resuming would run the node again and repeat its side effects.
-            return Flowable.<Event>error(
-                new UnsupportedOperationException(NODE_ROOT_RESUME_UNSUPPORTED));
-          }
           return resumeCore(session, resolvedInvocationId, newMessage, runConfig, stateDelta);
         });
   }
@@ -1067,50 +1067,52 @@ private Flowable<Event> resumeCore(
         () -> {
           Context capturedContext = Context.current();
 
-          // Not new input: Python drops it and its state delta rather than adding a second prompt.
+          // Not new input: Python drops newMessage, and outside the node path stateDelta too.
           boolean supersededByOriginalMessage =
               newMessage != null
                   && !carriesFunctionResponse(newMessage)
                   && findOriginalUserMessage(session, resolvedInvocationId).isPresent();
+          Content message = supersededByOriginalMessage ? null : newMessage;
+          Map<String, Object> delta =
+              supersededByOriginalMessage && node == null ? null : stateDelta;
 
           // Merged eagerly so the on-user-message callback sees it.
-          if (!supersededByOriginalMessage && stateDelta != null && !stateDelta.isEmpty()) {
-            stateDelta.forEach((key, value) -> session.state().put(key, value));
+          if (delta != null && !delta.isEmpty()) {
+            delta.forEach((key, value) -> session.state().put(key, value));
           }
 
           // For the pre-run plugin callbacks; the agent is resolved later, in runResumedAgent.
           InvocationContext initialContext =
               newInvocationContextBuilder(session)
                   .invocationId(resolvedInvocationId)
                   .runConfig(runConfig)
-                  .userContent(
-                      newMessage != null
-                          ? newMessage
-                          : originalUserMessage(session, resolvedInvocationId))
+                  .userContent(userContentFor(session, resolvedInvocationId, newMessage))
                   .build();
 
           Flowable<Event> events;
-          if (supersededByOriginalMessage) {
-            events =
-                runResumedAgent(session, resolvedInvocationId, /* userContent= */ null, runConfig);
-          } else if (newMessage != null) {
+          if (message != null) {
             // Inherit the branch of the call it answers, so routing and rehydration see it.
             String branch =
-                matchingFunctionCallEvent(session, newMessage).flatMap(Event::branch).orElse(null);
+                matchingFunctionCallEvent(session, message).flatMap(Event::branch).orElse(null);
             events =
                 this.pluginManager
-                    .onUserMessageCallback(initialContext, newMessage)
+                    .onUserMessageCallback(initialContext, message)
                     .compose(Tracing.<Content>withContext(capturedContext))
-                    .defaultIfEmpty(newMessage)
+                    .defaultIfEmpty(message)
                     .flatMap(
-                        content ->
-                            appendNewMessageToSession(
-                                session,
-                                content,
-                                initialContext,
-                                runConfig.saveInputBlobsAsArtifacts(),
-                                stateDelta,
-                                branch))
+                        content -> {
+                          // A plugin's message must not mix answers and text either.
+                          if (node != null) {
+                            checkNoTextWithFunctionResponses(content);
+                          }
+                          return appendNewMessageToSession(
+                              session,
+                              content,
+                              initialContext,
+                              runConfig.saveInputBlobsAsArtifacts(),
+                              delta,
+                              branch);
+                        })
                     // Persist first: an unpersisted response leaves the call looking unanswered.
                     .flatMap(userEvent -> this.sessionService.appendEvent(session, userEvent))
                     .flatMapPublisher(
@@ -1120,7 +1122,7 @@ private Flowable<Event> resumeCore(
                                 resolvedInvocationId,
                                 userEvent.content().orElse(null),
                                 runConfig));
-          } else if (stateDelta != null && !stateDelta.isEmpty()) {
+          } else if (delta != null && !delta.isEmpty()) {
             // No message to carry the delta, so persist a content-less event, as Python does.
             Event stateDeltaEvent =
                 Event.builder()
@@ -1129,9 +1131,7 @@ private Flowable<Event> resumeCore(
                     .invocationId(resolvedInvocationId)
                     .author(Role.USER)
                     .actions(
-                        EventActions.builder()
-                            .stateDelta(new ConcurrentHashMap<>(stateDelta))
-                            .build())
+                        EventActions.builder().stateDelta(new ConcurrentHashMap<>(delta)).build())
                     .build();
             events =
                 this.sessionService
@@ -1168,10 +1168,7 @@ private Flowable<Event> runResumedAgent(
               newInvocationContextBuilder(session, rootAgent)
                   .invocationId(resolvedInvocationId)
                   .runConfig(runConfig)
-                  .userContent(
-                      userContent != null
-                          ? userContent
-                          : originalUserMessage(session, resolvedInvocationId))
+                  .userContent(userContentFor(session, resolvedInvocationId, userContent))
                   .build();
           // Rehydrate before routing, as Python does, so a root that checkpointed re-enters itself.
           rootContext.populateInvocationAgentStates();
@@ -1340,18 +1337,24 @@ private static void checkNoTextWithFunctionResponses(@Nullable Content message)
   }
 
   /**
-   * For a node root, rejects a message whose function responses answer calls from an earlier
-   * invocation: without replay, running the node again would repeat its side effects. As in Python
-   * and Kotlin, responses without an id are ignored, so they start a new invocation.
+   * Returns {@code context} moved to the invocation whose calls {@code message} answers when the
+   * root is a workflow node, or unchanged otherwise. A plugin may replace the user's message with
+   * {@code message}, so it is checked again.
    *
-   * @throws IllegalArgumentException if a response id matches no call, or the ids span invocations.
-   * @throws UnsupportedOperationException if the responses answer an earlier invocation.
+   * @throws IllegalArgumentException if {@code message} mixes answers and text, or its answers do
+   *     not resolve to one invocation.
    */
-  private static void checkNoAnswerToAnEarlierInvocation(Session session, Content message) {
-    validateResumeMessage(session, message, /* requireFirstId= */ false);
-    if (matchingFunctionCallEvent(session, message).isPresent()) {
-      throw new UnsupportedOperationException(NODE_ROOT_RESUME_UNSUPPORTED);
+  private InvocationContext resolveNodeRootContext(
+      Session session, InvocationContext context, Content message) {
+    if (node == null) {
+      return context;
     }
+    checkNoTextWithFunctionResponses(message);
+    validateResumeMessage(session, message, /* requireFirstId= */ false);
+    String answeredInvocationId = resolveInvocationId(session, message, /* invocationId= */ null);
+    return answeredInvocationId == null
+        ? context
+        : context.toBuilder().invocationId(answeredInvocationId).build();
   }
 
   /**
@@ -1427,6 +1430,23 @@ private static Content originalUserMessage(Session session, String invocationId)
                     "No user message to resume the requested invocation from."));
   }
 
+  /**
+   * Returns the user content for a run of {@code invocationId}: {@code message}, or the message
+   * that started the invocation when there is none. As in Python, a workflow node root also gets
+   * that message when {@code message} answers calls.
+   *
+   * @throws IllegalArgumentException if there is no message and the invocation has no user message
+   */
+  private Content userContentFor(Session session, String invocationId, @Nullable Content message) {
+    if (message == null) {
+      return originalUserMessage(session, invocationId);
+    }
+    if (node != null && carriesFunctionResponse(message)) {
+      return findOriginalUserMessage(session, invocationId).orElse(message);
+    }
+    return message;
+  }
+
   private static boolean carriesFunctionResponse(Content message) {
     return message.parts().stream()
         .flatMap(List::stream)