diff --git a/core/src/main/java/com/google/adk/agents/Context.java b/core/src/main/java/com/google/adk/agents/Context.java index 0301b6079..37e2cf9d9 100644 --- a/core/src/main/java/com/google/adk/agents/Context.java +++ b/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 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 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}. diff --git a/core/src/main/java/com/google/adk/agents/LlmAgent.java b/core/src/main/java/com/google/adk/agents/LlmAgent.java index d36f0bda0..3bf5dbda6 100644 --- a/core/src/main/java/com/google/adk/agents/LlmAgent.java +++ b/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 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() { return executor; } diff --git a/core/src/main/java/com/google/adk/runner/Runner.java b/core/src/main/java/com/google/adk/runner/Runner.java index 96ae41f96..d80f60d4b 100644 --- a/core/src/main/java/com/google/adk/runner/Runner.java +++ b/core/src/main/java/com/google/adk/runner/Runner.java @@ -111,8 +111,6 @@ public class Runner { private final UuidProvider uuidProvider; private final ConcurrentMap 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 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 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 runAsync( RunConfig runConfig, @Nullable Map stateDelta) { checkState( - isResumable(), + isResumable() || node != null, "Resuming an invocation requires an App configured with a resumable ResumabilityConfig."); Flowable result = Flowable.defer( @@ -782,14 +781,18 @@ protected Flowable 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 runNewInvocation( Session session, Content newMessage, @@ -797,11 +800,6 @@ private Flowable runNewInvocation( @Nullable Map 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 runNewInvocation( .onUserMessageCallback(initialContext, newMessage) .compose(Tracing.withContext(capturedContext)) .defaultIfEmpty(newMessage) - .flatMap( - content -> - createUserMessageEvent( - session, - content, - initialContext, - runConfig.saveInputBlobsAsArtifacts(), - stateDelta)) .flatMapPublisher( - userEvent -> - runAgentForUserEvent(initialContext, session, userEvent, rootAgent) - .compose(Tracing.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.withContext(capturedContext))); + }) .doOnError(throwable -> runOnRunError(initialContext, throwable)); }) .doOnError( @@ -893,7 +894,11 @@ private Flowable 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 runResumableFromSession( Session session, @@ -1043,11 +1048,6 @@ private Flowable runResumableFromSession( return Flowable.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.error( - new UnsupportedOperationException(NODE_ROOT_RESUME_UNSUPPORTED)); - } return resumeCore(session, resolvedInvocationId, newMessage, runConfig, stateDelta); }); } @@ -1067,15 +1067,18 @@ private Flowable 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 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. @@ -1083,34 +1086,33 @@ private Flowable resumeCore( newInvocationContextBuilder(session) .invocationId(resolvedInvocationId) .runConfig(runConfig) - .userContent( - newMessage != null - ? newMessage - : originalUserMessage(session, resolvedInvocationId)) + .userContent(userContentFor(session, resolvedInvocationId, newMessage)) .build(); Flowable 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.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 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 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 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) diff --git a/core/src/main/java/com/google/adk/workflow/NodeExecutionState.java b/core/src/main/java/com/google/adk/workflow/NodeExecutionState.java index f1c53dab1..caf2aca6e 100644 --- a/core/src/main/java/com/google/adk/workflow/NodeExecutionState.java +++ b/core/src/main/java/com/google/adk/workflow/NodeExecutionState.java @@ -20,12 +20,17 @@ import static com.google.common.base.Preconditions.checkState; import com.google.adk.agents.Context; +import com.google.adk.agents.InvocationContext; import com.google.adk.annotations.Experimental; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import com.google.errorprone.annotations.CanIgnoreReturnValue; import java.util.Collection; +import java.util.Collections; +import java.util.LinkedHashMap; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; @@ -54,7 +59,10 @@ public final class NodeExecutionState { private final String runId; private final String nodePath; private final ImmutableList outputFor; + private final Map resumeInputs; private final ConcurrentMap childRunCounts = new ConcurrentHashMap<>(); + // The paths of the node runs started but not finished in this turn, shared under one root. + private final Set unfinishedNodePaths; private final Set interruptIds = new CopyOnWriteArraySet<>(); private final AtomicReference outputState = new AtomicReference<>(NO_OUTPUT); private String eventAuthor; @@ -62,6 +70,7 @@ public final class NodeExecutionState { private boolean routesEmitted; private @Nullable Exception failure; private @Nullable String failureNodePath; + private @Nullable ImmutableMap recoveredChildren; private NodeExecutionState(Builder builder) { checkState(builder.node != null, "A node state needs a node."); @@ -84,6 +93,9 @@ private NodeExecutionState(Builder builder) { ? ImmutableList.builder().add(nodePath).addAll(outputParent.outputFor).build() : ImmutableList.of(nodePath); this.eventAuthor = parentState == null ? "" : parentState.eventAuthor; + this.unfinishedNodePaths = + parentState == null ? ConcurrentHashMap.newKeySet() : parentState.unfinishedNodePaths; + this.resumeInputs = Collections.unmodifiableMap(new LinkedHashMap<>(builder.resumeInputs)); } /** Returns a builder for the state of one node run. */ @@ -181,6 +193,15 @@ public ImmutableSet interruptIds() { return ImmutableSet.copyOf(interruptIds); } + /** + * Returns the user's answers to the long-running calls that this node or its descendants raised + * in an earlier turn, keyed by interrupt ID. A null answer has a null value in the map. The map + * is empty unless the node reruns on resume. + */ + public Map resumeInputs() { + return resumeInputs; + } + NodeRunner.EventSink eventSink() { return eventSink; } @@ -193,11 +214,47 @@ ImmutableList outputFor() { return outputFor; } + /** Returns whether this run's output is also its parent's output, as with {@code useAsOutput}. */ + boolean outputCountsForParent() { + return outputParent != null; + } + /** Returns the next run ID for a child node named {@code name}, counting from {@code "1"}. */ String nextChildRunId(String name) { return String.valueOf(childRunCounts.merge(name, 1, Integer::sum)); } + /** Records that this run started; it counts as canceled mid-run until it finishes. */ + void markStarted() { + unfinishedNodePaths.add(nodePath); + } + + /** Records that this run finished, with or without a result. */ + void markFinished() { + unfinishedNodePaths.remove(nodePath); + } + + /** + * Returns whether this run's child {@code childKey}, a {@code name@runId} segment, started in + * this turn and has not finished. When read before the child runs again, this means a failing + * sibling or caller canceled it. + */ + boolean childCanceledMidRun(String childKey) { + return unfinishedNodePaths.contains(nodePath + "/" + childKey); + } + + /** + * Returns the states of this node's children from earlier turns, keyed by {@code name@runId}, + * which the first call reconstructs from the events of {@code invocationContext}. + */ + synchronized ImmutableMap recoveredChildren( + InvocationContext invocationContext) { + if (recoveredChildren == null) { + recoveredChildren = ResumeScan.scan(invocationContext, nodePath); + } + return recoveredChildren; + } + /** Returns whether an event has carried the node's own or delegated output. */ boolean outputEmitted() { OutputStatus status = outputState.get().status(); @@ -267,6 +324,20 @@ void releaseOutputDelegation() { outputState.compareAndSet(DELEGATED_OUTPUT, NO_OUTPUT); } + /** + * Records {@code output} from a {@code useAsOutput} child as this node's own unsent output when + * no event carried it for this run, as when the child replays an earlier turn's output. + * + * @throws IllegalStateException if the node has no pending delegation + */ + void adoptReplayedOutput(Object output) { + checkState( + outputState.compareAndSet( + DELEGATED_OUTPUT, new OutputState(OutputStatus.PRODUCED, checkNotNull(output))), + "Node '%s' adopted a replayed output without a pending delegation.", + node.name()); + } + /** Returns whether an event has carried the node's current route selection. */ boolean routesEmitted() { return routesEmitted; @@ -336,6 +407,7 @@ static final class Builder { private @Nullable Context parent; private @Nullable String runId; private boolean useAsOutput; + private Map resumeInputs = ImmutableMap.of(); private Builder() {} @@ -377,6 +449,13 @@ Builder useAsOutput(boolean useAsOutput) { return this; } + /** Sets the answers the node resumes with, keyed by interrupt ID. */ + @CanIgnoreReturnValue + Builder resumeInputs(Map resumeInputs) { + this.resumeInputs = checkNotNull(resumeInputs); + return this; + } + /** * Builds the state. * diff --git a/core/src/main/java/com/google/adk/workflow/NodeRunner.java b/core/src/main/java/com/google/adk/workflow/NodeRunner.java index 658b2d435..fd9ca1e77 100644 --- a/core/src/main/java/com/google/adk/workflow/NodeRunner.java +++ b/core/src/main/java/com/google/adk/workflow/NodeRunner.java @@ -27,7 +27,9 @@ import com.google.adk.annotations.Experimental; import com.google.adk.events.Event; import com.google.adk.events.EventActions; +import com.google.adk.workflow.ResumeScan.Interception; import com.google.common.base.CharMatcher; +import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import com.google.errorprone.annotations.CanIgnoreReturnValue; import com.google.genai.types.FinishReason; @@ -78,6 +80,7 @@ interface EventSink { private final boolean useAsOutput; private final boolean useSubBranch; private final @Nullable String overrideBranch; + private final Map resumeInputs; private NodeRunner(Builder builder) { checkState(builder.node != null, "A node runner needs a node."); @@ -100,6 +103,7 @@ private NodeRunner(Builder builder) { this.useAsOutput = builder.useAsOutput; this.useSubBranch = builder.useSubBranch; this.overrideBranch = builder.overrideBranch; + this.resumeInputs = builder.resumeInputs; } /** Returns a builder for a runner of one node. */ @@ -213,16 +217,45 @@ private static Maybe startChild( if (delegatesOutput) { parentState.claimOutputDelegation(); } + String runId = config.runId().orElseGet(() -> parentState.nextChildRunId(node.name())); + // Deferred so that a throw from scan or intercept releases the delegation, as a failure does. + return Single.defer(() -> runOrReplayChild(parent, node, nodeInput, runId, config)) + .doOnError( + error -> { + if (delegatesOutput) { + parentState.releaseOutputDelegation(); + } + }) + .flatMapMaybe(child -> childResult(parent, child, delegatesOutput, config.raiseOnWait())); + } + + /** + * Runs {@code node} as a child of {@code parent}, or replays the result it recorded in an earlier + * turn under the same run ID. + */ + private static Single runOrReplayChild( + Context parent, Node node, @Nullable Object nodeInput, String runId, RunNodeConfig config) { + NodeExecutionState parentState = parent.requireNodeState(); + String childKey = segment(node.name(), runId); + ResumeScan.RecoveredNode recovered = + parentState.recoveredChildren(parent.invocationContext()).get(childKey); + Interception interception = + ResumeScan.intercept( + node, recovered, /* dynamic= */ true, parentState.childCanceledMidRun(childKey)); + if (!interception.shouldRun()) { + return Single.just( + ResumeScan.replayContext(parent, node, runId, config.useAsOutput(), interception)); + } return builder() .node(node) .parent(parent) - .runId(config.runId().orElseGet(() -> parentState.nextChildRunId(node.name()))) + .runId(runId) .useAsOutput(config.useAsOutput()) .useSubBranch(config.useSubBranch()) .overrideBranch(config.overrideBranch().orElse(null)) + .resumeInputs(interception.resumeInputs()) .build() - .run(nodeInput) - .flatMapMaybe(child -> childResult(parent, child, delegatesOutput, config.raiseOnWait())); + .run(nodeInput); } /** @@ -259,7 +292,12 @@ private static Maybe childResult( && (childNode instanceof Workflow || childNode.waitForOutput())) { return Maybe.error(new NodeInterruptedException()); } - return Maybe.fromOptional(childState.output()); + // No event of this run carries a delegate's replayed output, so the caller emits it. + Optional output = childState.output(); + if (delegatesOutput && output.isPresent() && !childState.outputEmitted()) { + parentState.adoptReplayedOutput(output.get()); + } + return Maybe.fromOptional(output); } /** @@ -270,20 +308,40 @@ private static Maybe childResult( Single run(@Nullable Object nodeInput) { return Single.defer( () -> { - Context context = - new Context( - childInvocationContext(), - NodeExecutionState.builder() - .node(node) - .eventSink(eventSink) - .parent(parent) - .runId(runId) - .useAsOutput(useAsOutput) - .build()); + if (parent != null) { + Context context = newContext(resumeInputs); + return new Activation(context).run(nodeInput).toSingleDefault(context); + } + // A root scans its own progress and resumes like a dynamic child. + Interception interception = + ResumeScan.intercept( + node, + ResumeScan.scan(invocationContext, /* basePath= */ "") + .get(segment(node.name(), runId)), + /* dynamic= */ true, + /* canceledMidRun= */ false); + Context context = newContext(interception.resumeInputs()); + if (!interception.shouldRun()) { + interception.replayOn(context); + return new Activation(context).flushPending().toSingleDefault(context); + } return new Activation(context).run(nodeInput).toSingleDefault(context); }); } + private Context newContext(Map resumeInputs) { + return new Context( + childInvocationContext(), + NodeExecutionState.builder() + .node(node) + .eventSink(eventSink) + .parent(parent) + .runId(runId) + .useAsOutput(useAsOutput) + .resumeInputs(resumeInputs) + .build()); + } + private InvocationContext childInvocationContext() { if (useSubBranch) { String base = @@ -319,7 +377,10 @@ Completable run(@Nullable Object nodeInput) { // Avoid concatMapCompletable: its prefetch would let the node run ahead of the sink. .flatMapCompletable(this::dispatch, /* delayErrors= */ true, /* maxConcurrency= */ 1) .andThen(Completable.defer(this::flushPending)) - .onErrorResumeNext(this::onError); + .onErrorResumeNext(this::onError) + // A run cut short stays unfinished, so a later call in this turn reruns it. + .doOnSubscribe(unused -> nodeState.markStarted()) + .doOnComplete(nodeState::markFinished); } private Completable onError(Throwable error) { @@ -480,10 +541,15 @@ private record QueuedEvent( /** Appends a {@code name@runId} segment to {@code base}, joined by {@code separator}. */ static String appendSegment(@Nullable String base, String name, String runId, char separator) { - String segment = name + "@" + runId; + String segment = segment(name, runId); return isNullOrEmpty(base) ? segment : base + separator + segment; } + /** Returns the {@code name@runId} path segment of a node's run. */ + static String segment(String name, String runId) { + return name + "@" + runId; + } + /** Returns whether {@code event} is a final event whose content is the node's output. */ private static boolean isMessageAsOutput(Event event) { return !event.partial().orElse(false) @@ -524,6 +590,7 @@ static final class Builder { private boolean useAsOutput; private boolean useSubBranch; private @Nullable String overrideBranch; + private Map resumeInputs = ImmutableMap.of(); private Builder() {} @@ -589,6 +656,16 @@ Builder overrideBranch(@Nullable String overrideBranch) { return this; } + /** + * Sets the answers that a child resumes with, keyed by interrupt ID. A runner at the root reads + * the root's answers from the session itself. + */ + @CanIgnoreReturnValue + Builder resumeInputs(Map resumeInputs) { + this.resumeInputs = checkNotNull(resumeInputs); + return this; + } + /** * Builds the runner. * diff --git a/core/src/main/java/com/google/adk/workflow/NodeState.java b/core/src/main/java/com/google/adk/workflow/NodeState.java index ea6fb867f..58760274e 100644 --- a/core/src/main/java/com/google/adk/workflow/NodeState.java +++ b/core/src/main/java/com/google/adk/workflow/NodeState.java @@ -18,11 +18,15 @@ import com.google.adk.annotations.Experimental; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; /** A node's progress within one run of a workflow graph. */ @Experimental final class NodeState { + /** The key under which a workflow's checkpoint holds the checkpoints of its nodes. */ + static final String NODES_KEY = "nodes"; + /** Where the node stands in the current run. */ NodeStatus status; @@ -32,4 +36,13 @@ final class NodeState { NodeState(NodeStatus status) { this.status = status; } + + /** + * Returns what a resumable session records for the node. Its keys and numeric status are shared + * with the other ADK languages. The node's answers are left out: they can hold a credential the + * user sent, and a resume reads them from the session's events. + */ + ImmutableMap toCheckpoint() { + return ImmutableMap.of("status", status.code(), "interrupts", interrupts); + } } diff --git a/core/src/main/java/com/google/adk/workflow/NodeStatus.java b/core/src/main/java/com/google/adk/workflow/NodeStatus.java index f8cf43188..09fa98c55 100644 --- a/core/src/main/java/com/google/adk/workflow/NodeStatus.java +++ b/core/src/main/java/com/google/adk/workflow/NodeStatus.java @@ -18,15 +18,29 @@ import com.google.adk.annotations.Experimental; -/** The status of a node within one run of a workflow graph. */ +/** + * The status of a node within one run of a workflow graph. A resumable session records its {@link + * #code}, which the other ADK languages share, so the codes must not change. + */ @Experimental enum NodeStatus { /** Running now. */ - RUNNING, + RUNNING(2), /** Finished; a node that waits for output finishes only once it produces an output or routes. */ - COMPLETED, + COMPLETED(3), /** Paused, either awaiting a user's answers or awaiting another trigger. */ - WAITING + WAITING(4); + + private final int code; + + NodeStatus(int code) { + this.code = code; + } + + /** Returns the number that a resumable session records for this status. */ + int code() { + return code; + } } diff --git a/core/src/main/java/com/google/adk/workflow/RunNodeConfig.java b/core/src/main/java/com/google/adk/workflow/RunNodeConfig.java index 4a61e1c9e..dbabfea66 100644 --- a/core/src/main/java/com/google/adk/workflow/RunNodeConfig.java +++ b/core/src/main/java/com/google/adk/workflow/RunNodeConfig.java @@ -30,8 +30,10 @@ public abstract class RunNodeConfig { /** - * Returns the ID of the child's run, which names the run in node paths and branches. When empty, - * the child receives the next number among the caller's runs of nodes with the same name. + * Returns the ID of the child's run, which names the run in node paths and branches. A child that + * recorded an output under the same ID in an earlier turn returns it without running again. When + * this ID is empty, the child receives the next number among the caller's runs of nodes with the + * same name, in call order, so concurrent children need explicit IDs. */ public abstract Optional runId(); diff --git a/core/src/main/java/com/google/adk/workflow/Scheduler.java b/core/src/main/java/com/google/adk/workflow/Scheduler.java index 2f77f884b..08e7737f1 100644 --- a/core/src/main/java/com/google/adk/workflow/Scheduler.java +++ b/core/src/main/java/com/google/adk/workflow/Scheduler.java @@ -20,9 +20,15 @@ import static com.google.common.collect.ImmutableList.toImmutableList; import com.google.adk.agents.Context; +import com.google.adk.agents.InvocationContext; import com.google.adk.annotations.Experimental; +import com.google.adk.events.Event; +import com.google.adk.events.EventActions; +import com.google.adk.workflow.ResumeScan.Interception; +import com.google.adk.workflow.ResumeScan.RecoveredNode; import com.google.common.base.Splitter; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import io.reactivex.rxjava3.core.Completable; import io.reactivex.rxjava3.disposables.CompositeDisposable; @@ -32,6 +38,7 @@ import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; +import java.util.HashSet; import java.util.LinkedHashMap; import java.util.LinkedHashSet; import java.util.List; @@ -50,18 +57,25 @@ final class Scheduler { private final Graph graph; private final Context context; private final NodeExecutionState workflowState; - private final Map nodeStates = new HashMap<>(); + private final Map nodeStates = new LinkedHashMap<>(); private final Map nodeOutputs = new HashMap<>(); // The branch each completed node ran on, so a join can re-merge parallel branches. private final Map nodeBranches = new HashMap<>(); private final TriggerQueue triggerQueue = new TriggerQueue(); private final Set interruptIds = new LinkedHashSet<>(); + // The states of the workflow's nodes from earlier turns, keyed by name@runId. + private final ImmutableMap recovered; + // The nodes replayed from history whose completion is not handled yet. + private final Set replayed = new HashSet<>(); + // The nodes whose latest output was replayed, so no event in this turn carried it in outputFor. + private final Set replayedOutputs = new HashSet<>(); Scheduler(Workflow workflow, Graph graph, Context context) { this.workflow = workflow; this.graph = graph; this.context = context; this.workflowState = context.requireNodeState(); + this.recovered = ResumeScan.scan(context.invocationContext(), workflowState.nodePath()); } /** Runs the graph from its start node, which hands {@code nodeInput} to its successors. */ @@ -188,16 +202,54 @@ private void scheduleReady() { private void start(String nodeName, Trigger trigger) { String runId = workflowState.nextChildRunId(nodeName); nodeStates.put(nodeName, new NodeState(NodeStatus.RUNNING)); + Node node = graph.node(nodeName); + boolean useAsOutput = graph.terminalNodeNames().contains(nodeName); + String childKey = NodeRunner.segment(nodeName, runId); + RecoveredNode recoveredNode = recovered.get(childKey); + Interception interception = + ResumeScan.intercept( + node, + recoveredNode, + /* dynamic= */ false, + workflowState.childCanceledMidRun(childKey)); + if (recoveredNode != null && !interception.shouldRun()) { + // Only a replay that completes an answered transfer marks a new step, as a run does. + boolean completesTransfer = + interception.transferToAgent() != null + && !recoveredNode.interruptIds.isEmpty() + && recoveredNode.unresolved().isEmpty(); + if (completesTransfer) { + emitCheckpoint(); + } else { + replayed.add(nodeName); + } + Context replayContext = + ResumeScan.replayContext(context, node, runId, useAsOutput, interception); + // A replay counts as running until its completion is handled, as a run does. + running.put(nodeName, Disposable.empty()); + completions.add(new Completion(nodeName, replayContext)); + wakeUps.onNext(true); + return; + } + Completable checkpointed = checkpoint(); + // Processing the checkpoint can end the run, as when the caller stops after that event. + if (childRuns.isDisposed()) { + return; + } + // The node starts once its checkpoint is processed, as each of its own events is. Disposable run = - NodeRunner.builder() - .node(graph.node(nodeName)) - .parent(context) - .runId(runId) - .useAsOutput(graph.terminalNodeNames().contains(nodeName)) - .useSubBranch(trigger.useSubBranch()) - .overrideBranch(trigger.branch()) - .build() - .run(trigger.input()) + checkpointed + .andThen( + NodeRunner.builder() + .node(node) + .parent(context) + .runId(runId) + .useAsOutput(useAsOutput) + .useSubBranch(trigger.useSubBranch()) + .overrideBranch(trigger.branch()) + .resumeInputs(interception.resumeInputs()) + .build() + .run(trigger.input())) .subscribe( childContext -> { completions.add(new Completion(nodeName, childContext)); @@ -213,25 +265,97 @@ private void start(String nodeName, Trigger trigger) { private void handleCompletion(String nodeName, Context childContext) { NodeState progress = nodeStates.get(nodeName); NodeExecutionState childState = childContext.requireNodeState(); + boolean wasReplayed = replayed.remove(nodeName); ImmutableSet childInterruptIds = childState.interruptIds(); if (!childInterruptIds.isEmpty()) { progress.status = NodeStatus.WAITING; progress.interrupts = childInterruptIds.asList(); interruptIds.addAll(childInterruptIds); + emitCheckpoint(); return; } if (graph.node(nodeName).waitForOutput() && childState.output().isEmpty() && childState.selectedRoutes().isEmpty()) { progress.status = NodeStatus.WAITING; + emitCheckpoint(); return; } progress.status = NodeStatus.COMPLETED; nodeBranches.put(nodeName, childContext.invocationContext().branch().orElse(null)); - childState.output().ifPresent(output -> nodeOutputs.put(nodeName, output)); + Optional output = childState.output(); + if (output.isPresent()) { + nodeOutputs.put(nodeName, output.get()); + if (wasReplayed) { + replayedOutputs.add(nodeName); + } else { + replayedOutputs.remove(nodeName); + } + } + if (wasReplayed) { + reemitReplayedOutput(childContext); + } else { + emitCheckpoint(); + } enqueueSuccessorTriggers(nodeName, childContext); } + /** Queues a checkpoint, as {@link #checkpoint} does, without waiting for it to be processed. */ + private void emitCheckpoint() { + Completable unused = checkpoint(); + } + + /** + * On a resumable invocation, records where every node stands, so that a later turn can see how + * far this one got. The event is queued before any event that follows it, and the returned + * Completable completes once it is processed. + */ + private Completable checkpoint() { + if (!context.invocationContext().isResumable()) { + return Completable.complete(); + } + Map nodes = new LinkedHashMap<>(); + nodeStates.forEach((name, progress) -> nodes.put(name, progress.toCheckpoint())); + return send( + workflowEvent(context.invocationContext().branch().orElse(null)) + .actions( + EventActions.builder() + .agentState(ImmutableMap.of(NodeState.NODES_KEY, nodes)) + .build())); + } + + /** + * On a resumable invocation, re-emits the output of a node replayed from history on the node's + * branch, without the paths that the output counts for. + */ + private void reemitReplayedOutput(Context childContext) { + Optional output = childContext.requireNodeState().output(); + if (!context.invocationContext().isResumable() || output.isEmpty()) { + return; + } + Completable unused = + send( + workflowEvent(childContext.invocationContext().branch().orElse(null)) + .output(output.get()) + .nodeInfo( + NodeInfo.builder().path(childContext.requireNodeState().nodePath()).build())); + } + + /** Returns the start of an event that the workflow itself emits on {@code branch}. */ + private Event.Builder workflowEvent(@Nullable String branch) { + InvocationContext invocationContext = context.invocationContext(); + return Event.builder() + .id(invocationContext.newUuid()) + .timestamp(invocationContext.now().toEpochMilli()) + .invocationId(invocationContext.invocationId()) + .author(workflow.name()) + .branch(branch); + } + + private Completable send(Event.Builder event) { + return workflowState.eventSink().send(event.build(), /* nodeFailure= */ null); + } + private void enqueueSuccessorTriggers(String nodeName, Context childContext) { NodeExecutionState childState = childContext.requireNodeState(); ImmutableList successors = @@ -306,11 +430,20 @@ private Completable finalizeRun() { workflow.name(), terminalOutputs.size()))); } if (terminalOutputs.size() == 1) { - workflowState.produceOutput(nodeOutputs.get(terminalOutputs.get(0))); - // The terminal node's event already carried the output for the workflow as well. - workflowState.markOutputEmitted(); + String terminalNode = terminalOutputs.get(0); + workflowState.produceOutput(nodeOutputs.get(terminalNode)); + // A replayed terminal node did not carry the output for a delegating parent. + if (!replayedOutputs.contains(terminalNode) || !workflowState.outputCountsForParent()) { + workflowState.markOutputEmitted(); + } + } + // On a resumable invocation, a clean finish is marked so a later turn can see it. + if (!context.invocationContext().isResumable()) { + return Completable.complete(); } - return Completable.complete(); + return send( + workflowEvent(context.invocationContext().branch().orElse(null)) + .actions(EventActions.builder().endOfAgent(true).build())); } /** Returns the longest segment prefix that the dot-separated {@code branches} share, if any. */ diff --git a/core/src/test/java/com/google/adk/runner/RunnerRootNodeTest.java b/core/src/test/java/com/google/adk/runner/RunnerRootNodeTest.java index 801f251b3..6c922f2b3 100644 --- a/core/src/test/java/com/google/adk/runner/RunnerRootNodeTest.java +++ b/core/src/test/java/com/google/adk/runner/RunnerRootNodeTest.java @@ -53,6 +53,7 @@ import com.google.adk.workflow.Workflow; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; import com.google.genai.types.Content; import com.google.genai.types.FinishReason; import com.google.genai.types.FunctionCall; @@ -368,6 +369,42 @@ public Maybe onUserMessageCallback( assertThat(inputs).containsExactly(rewritten); } + @Test + public void runAsync_resumingANodeRootWithAStateDelta_showsItToTheOnUserMessageCallback() { + List> seen = new ArrayList<>(); + Plugin stateReader = + new BasePlugin("state_reader") { + @Override + public Maybe onUserMessageCallback( + InvocationContext invocationContext, Content userMessage) { + seen.add(Optional.ofNullable(invocationContext.session().state().get("step"))); + return Maybe.empty(); + } + }; + Node ask = + new TestNode("ask", (context, input) -> Flowable.just(longRunningCallEvent("call_1"))); + Runner runner = runner(ask, stateReader); + Session session = newSession(runner); + runner.runAsync(USER_ID, session.id(), USER_MESSAGE).blockingSubscribe(); + + runner + .runAsync( + USER_ID, + session.id(), + /* invocationId= */ null, + functionResponse( + FunctionResponse.builder() + .id("call_1") + .name("tool") + .response(ImmutableMap.of("v", 1)) + .build()), + RunConfig.builder().build(), + ImmutableMap.of("step", "two")) + .blockingSubscribe(); + + assertThat(seen).containsExactly(Optional.empty(), Optional.of("two")).inOrder(); + } + @Test public void runAsync_partialEventOfANodeRoot_isStreamedButNotPersisted() { Node node = @@ -414,7 +451,7 @@ public void runAsync_takingTheFirstEventOfANodeRoot_cancelsTheNode() { } @Test - public void runAsync_resumingANodeRootInvocation_failsWithoutRunningTheNodeAgain() { + public void runAsync_resumingAFinishedNodeRootInvocation_replaysItsOutputWithoutRunningIt() { AtomicInteger runs = new AtomicInteger(); Node node = new TestNode( @@ -428,7 +465,7 @@ public void runAsync_resumingANodeRootInvocation_failsWithoutRunningTheNodeAgain String invocationId = runner.runAsync(USER_ID, session.id(), USER_MESSAGE).blockingFirst().invocationId(); - TestSubscriber subscriber = + List events = runner .runAsync( USER_ID, @@ -437,11 +474,61 @@ public void runAsync_resumingANodeRootInvocation_failsWithoutRunningTheNodeAgain /* newMessage= */ null, RunConfig.builder().build(), /* stateDelta= */ null) - .test(); - subscriber.awaitDone(5, SECONDS); + .toList() + .blockingGet(); - subscriber.assertError(UnsupportedOperationException.class); assertThat(runs.get()).isEqualTo(1); + Event replayed = getOnlyElement(events); + assertThat(replayed.output()).hasValue("done"); + assertThat(replayed.invocationId()).isEqualTo(invocationId); + } + + @Test + public void runAsync_resumingANodeRootWithTextWithoutResumability_keepsOnlyTheStateDelta() { + List inputs = new ArrayList<>(); + Runner runner = runner(inputRecorder(inputs)); + Session session = newSession(runner); + runner.runAsync(USER_ID, session.id(), USER_MESSAGE).blockingSubscribe(); + String invocationId = getOnlyElement(reload(runner, session).immutableEvents()).invocationId(); + + runner + .runAsync( + USER_ID, + session.id(), + invocationId, + Content.fromParts(Part.fromText("more")), + RunConfig.builder().build(), + ImmutableMap.of("step", "two")) + .blockingSubscribe(); + + // As in Python, the node reruns on the message that started the invocation. + assertThat(inputs).containsExactly(USER_MESSAGE, USER_MESSAGE); + Session reloaded = reload(runner, session); + assertThat(texts(reloaded.immutableEvents())).doesNotContain("more"); + assertThat(reloaded.state()).containsEntry("step", "two"); + assertThat(reloaded.immutableEvents().stream().map(Event::invocationId).distinct()) + .containsExactly(invocationId); + } + + @Test + public void runAsync_resumingANodeRootWithoutResumabilityOrAMessage_rerunsOnItsFirstMessage() { + List inputs = new ArrayList<>(); + Runner runner = runner(inputRecorder(inputs)); + Session session = newSession(runner); + runner.runAsync(USER_ID, session.id(), USER_MESSAGE).blockingSubscribe(); + String invocationId = getOnlyElement(reload(runner, session).immutableEvents()).invocationId(); + + runner + .runAsync( + USER_ID, + session.id(), + invocationId, + /* newMessage= */ null, + RunConfig.builder().build(), + /* stateDelta= */ null) + .blockingSubscribe(); + + assertThat(inputs).containsExactly(USER_MESSAGE, USER_MESSAGE); } @Test @@ -489,26 +576,156 @@ public void runAsync_textWithAFunctionResponseWithoutIdToANodeRoot_runsTheNode() } @Test - public void runAsync_functionResponseToTheNodeRootsEarlierCall_failsWithoutRunningTheNodeAgain() { - AtomicInteger runs = new AtomicInteger(); + public void runAsync_functionResponseToTheNodeRootsEarlierCall_runsTheNodeInThatInvocation() { + List inputs = new ArrayList<>(); Node node = new TestNode( "solo", (context, input) -> { - runs.incrementAndGet(); + inputs.add(input); return Flowable.just(functionCallEvent(IDENTIFIED_RESPONSE.id().orElseThrow())); }); Runner runner = runner(node); Session session = newSession(runner); + String invocationId = + runner.runAsync(USER_ID, session.id(), USER_MESSAGE).blockingFirst().invocationId(); + + runner + .runAsync(USER_ID, session.id(), functionResponse(IDENTIFIED_RESPONSE)) + .blockingSubscribe(); + + // The call was not long-running, so the node reruns on the invocation's first message. + assertThat(inputs).containsExactly(USER_MESSAGE, USER_MESSAGE); + assertThat( + reload(runner, session).immutableEvents().stream().map(Event::invocationId).distinct()) + .containsExactly(invocationId); + } + + @Test + public void runAsync_pluginReplacingTheMessageWithAnAnswer_completesTheNodeInThatInvocation() { + Node node = + new TestNode("ask", (context, input) -> Flowable.just(longRunningCallEvent("ask_1"))); + Content answer = + functionResponse( + FunctionResponse.builder() + .id("ask_1") + .name("tool") + .response(ImmutableMap.of("result", "yes")) + .build()); + AtomicBoolean replace = new AtomicBoolean(); + Plugin answeringPlugin = + new BasePlugin("answering") { + @Override + public Maybe onUserMessageCallback( + InvocationContext invocationContext, Content userMessage) { + return replace.get() ? Maybe.just(answer) : Maybe.empty(); + } + }; + Runner runner = runner(node, answeringPlugin); + Session session = newSession(runner); + String invocationId = + runner.runAsync(USER_ID, session.id(), USER_MESSAGE).blockingFirst().invocationId(); + replace.set(true); + + List events = + runner + .runAsync(USER_ID, session.id(), Content.fromParts(Part.fromText("ignored"))) + .toList() + .blockingGet(); + + Event output = getOnlyElement(events); + assertThat(output.output()).hasValue("yes"); + assertThat(output.invocationId()).isEqualTo(invocationId); + } + + @Test + public void runAsync_pluginReplacingTheMessageWithTextAndAnAnswer_failsBeforeRunningTheNode() { + AtomicInteger runs = new AtomicInteger(); + Plugin mixingPlugin = + new BasePlugin("mixing") { + @Override + public Maybe onUserMessageCallback( + InvocationContext invocationContext, Content userMessage) { + return Maybe.just(textWithFunctionResponse(IDENTIFIED_RESPONSE)); + } + }; + Runner runner = runner(countingNode(runs), mixingPlugin); + Session session = newSession(runner); + Flowable events = runner.runAsync(USER_ID, session.id(), USER_MESSAGE); + + IllegalArgumentException error = + assertThrows(IllegalArgumentException.class, events::blockingSubscribe); + + assertThat(error).hasMessageThat().contains("both function responses and text"); + assertThat(runs.get()).isEqualTo(0); + assertThat(reload(runner, session).immutableEvents()).isEmpty(); + } + + @Test + public void runAsync_pluginReplacingTheMessageWithAnAnswerToAnUnknownCall_failsBeforeRunning() { + AtomicInteger runs = new AtomicInteger(); + Plugin answeringPlugin = + new BasePlugin("answering") { + @Override + public Maybe onUserMessageCallback( + InvocationContext invocationContext, Content userMessage) { + return Maybe.just(functionResponse(IDENTIFIED_RESPONSE)); + } + }; + Runner runner = runner(countingNode(runs), answeringPlugin); + Session session = newSession(runner); + Flowable events = runner.runAsync(USER_ID, session.id(), USER_MESSAGE); + + IllegalArgumentException error = + assertThrows(IllegalArgumentException.class, events::blockingSubscribe); + + assertThat(error).hasMessageThat().contains("No matching function call"); + assertThat(runs.get()).isEqualTo(0); + } + + @Test + public void runAsync_pluginReplacingAResumableNodeRootsAnswerWithTextAndAnAnswer_fails() { + AtomicInteger runs = new AtomicInteger(); + Node node = + new TestNode( + "ask", + (context, input) -> { + runs.incrementAndGet(); + return Flowable.just(longRunningCallEvent("ask_1")); + }); + FunctionResponse answer = + FunctionResponse.builder() + .id("ask_1") + .name("tool") + .response(ImmutableMap.of("result", "yes")) + .build(); + AtomicBoolean replace = new AtomicBoolean(); + Plugin mixingPlugin = + new BasePlugin("mixing") { + @Override + public Maybe onUserMessageCallback( + InvocationContext invocationContext, Content userMessage) { + return replace.get() ? Maybe.just(textWithFunctionResponse(answer)) : Maybe.empty(); + } + }; + Runner runner = + runner( + App.builder() + .name(APP_NAME) + .rootNode(node) + .plugins(mixingPlugin) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build()); + Session session = newSession(runner); runner.runAsync(USER_ID, session.id(), USER_MESSAGE).blockingSubscribe(); - int eventsAfterTheFirstRun = reload(runner, session).immutableEvents().size(); - Flowable events = - runner.runAsync(USER_ID, session.id(), functionResponse(IDENTIFIED_RESPONSE)); + replace.set(true); + Flowable events = runner.runAsync(USER_ID, session.id(), functionResponse(answer)); - assertThrows(UnsupportedOperationException.class, events::blockingSubscribe); + IllegalArgumentException error = + assertThrows(IllegalArgumentException.class, events::blockingSubscribe); + assertThat(error).hasMessageThat().contains("both function responses and text"); assertThat(runs.get()).isEqualTo(1); - assertThat(reload(runner, session).immutableEvents()).hasSize(eventsAfterTheFirstRun); } @Test @@ -659,6 +876,16 @@ private static Node failingNode(Exception failure) { }); } + /** Returns a node that adds each input it runs on to {@code inputs} and emits nothing. */ + private static Node inputRecorder(List inputs) { + return new TestNode( + "solo", + (context, input) -> { + inputs.add(input); + return Flowable.empty(); + }); + } + private static Node countingNode(AtomicInteger runs) { return new TestNode( "solo", @@ -702,6 +929,13 @@ private static Event functionCallEvent(String callId) { .build(); } + /** Returns an event in which a node waits on the long-running call {@code callId}. */ + private static Event longRunningCallEvent(String callId) { + return functionCallEvent(callId).toBuilder() + .longRunningToolIds(ImmutableSet.of(callId)) + .build(); + } + private static ImmutableList ids(List events) { return events.stream().map(Event::id).collect(toImmutableList()); } diff --git a/core/src/test/java/com/google/adk/workflow/WorkflowResumeTest.java b/core/src/test/java/com/google/adk/workflow/WorkflowResumeTest.java new file mode 100644 index 000000000..1a32ef393 --- /dev/null +++ b/core/src/test/java/com/google/adk/workflow/WorkflowResumeTest.java @@ -0,0 +1,1111 @@ +/* + * Copyright 2026 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.workflow; + +import static com.google.adk.testing.WorkflowTestNodes.dispatcher; +import static com.google.adk.testing.WorkflowTestNodes.emitter; +import static com.google.adk.testing.WorkflowTestNodes.workflow; +import static com.google.adk.workflow.WorkflowResumeTesting.AGE_SCHEMA; +import static com.google.adk.workflow.WorkflowResumeTesting.LONG_RUNNING_CALL; +import static com.google.adk.workflow.WorkflowResumeTesting.START; +import static com.google.adk.workflow.WorkflowResumeTesting.answer; +import static com.google.adk.workflow.WorkflowResumeTesting.callContent; +import static com.google.adk.workflow.WorkflowResumeTesting.captureInput; +import static com.google.adk.workflow.WorkflowResumeTesting.go; +import static com.google.adk.workflow.WorkflowResumeTesting.modelText; +import static com.google.adk.workflow.WorkflowResumeTesting.newEvent; +import static com.google.adk.workflow.WorkflowResumeTesting.outputEventAt; +import static com.google.adk.workflow.WorkflowResumeTesting.pathOf; +import static com.google.adk.workflow.WorkflowResumeTesting.pathsOf; +import static com.google.adk.workflow.WorkflowResumeTesting.replyInput; +import static com.google.common.collect.ImmutableList.toImmutableList; +import static com.google.common.collect.MoreCollectors.onlyElement; +import static com.google.common.truth.Truth.assertThat; + +import com.google.adk.agents.Context; +import com.google.adk.agents.InvocationContext; +import com.google.adk.agents.LlmAgent; +import com.google.adk.events.Event; +import com.google.adk.events.EventActions; +import com.google.adk.models.LlmRequest; +import com.google.adk.plugins.BasePlugin; +import com.google.adk.plugins.Plugin; +import com.google.adk.testing.TestLlm; +import com.google.adk.testing.TestNode; +import com.google.adk.testing.TestUtils; +import com.google.adk.tools.BaseTool; +import com.google.adk.tools.ToolContext; +import com.google.adk.workflow.WorkflowResumeTesting.FailedTurn; +import com.google.adk.workflow.WorkflowResumeTesting.MessageOutputNode; +import com.google.adk.workflow.WorkflowResumeTesting.WaitingNode; +import com.google.adk.workflow.WorkflowResumeTesting.WorkflowSession; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; +import com.google.common.collect.Iterables; +import com.google.common.collect.Streams; +import com.google.genai.types.Content; +import com.google.genai.types.FunctionDeclaration; +import com.google.genai.types.Part; +import com.google.genai.types.Schema; +import com.google.genai.types.Type; +import io.reactivex.rxjava3.core.Completable; +import io.reactivex.rxjava3.core.Flowable; +import io.reactivex.rxjava3.core.Maybe; +import io.reactivex.rxjava3.core.Single; +import io.reactivex.rxjava3.schedulers.Schedulers; +import io.reactivex.rxjava3.subjects.CompletableSubject; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicInteger; +import org.jspecify.annotations.Nullable; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * Tests for pausing a workflow until the user answers and resuming it on a later turn, which + * rebuilds each node's progress from the invocation's session events. + */ +@RunWith(JUnit4.class) +public final class WorkflowResumeTest { + private static final RunNodeConfig USE_AS_OUTPUT = + RunNodeConfig.builder().runId("c").useAsOutput(true).build(); + + @Test + public void resume_withoutResumability_completesTheNodeAndRunsItsSuccessor() { + Node ask = longRunningAsk("ask", "ask_1"); + WorkflowSession session = + new WorkflowSession( + workflow(new Edge(START, ask), new Edge(ask, captureInput("reply"))), + /* resumable= */ false); + session.advance(go()); + + ImmutableList events = session.turn(answer("ask_1", ImmutableMap.of("v", 1))); + + assertThat(replyInput(events)).isEqualTo(ImmutableMap.of("v", 1)); + assertThat(events.stream().map(Event::invocationId).distinct()).hasSize(1); + assertThat(session.invocationIds()).hasSize(1); + } + + @Test + public void resume_nodeCompletedBeforeTheInterrupt_replaysItsOutputWithoutRunning() { + CountingNode prep = new CountingNode("prep", "prep_out"); + Node ask = longRunningAsk("ask", "ask_1"); + WorkflowSession session = + new WorkflowSession( + workflow( + new Edge(START, prep), new Edge(prep, ask), new Edge(ask, captureInput("reply")))); + session.advance(go()); + + ImmutableList events = session.turn(answer("ask_1", ImmutableMap.of("v", 1))); + + assertThat(prep.runs()).isEqualTo(1); + // The replayed output is emitted again but counts for no node, since prep did not run. + Event prepEvent = outputEventAt(events, "wf@1/prep@1"); + assertThat(prepEvent.output()).hasValue("prep_out"); + assertThat(prepEvent.nodeInfo().orElseThrow().outputFor()).isEmpty(); + } + + @Test + public void resume_routingNodeCompletedBeforeTheInterrupt_replaysItsRouteWithoutRunning() { + AtomicInteger routerRuns = new AtomicInteger(); + Route yes = new Route.Tag("yes"); + Node router = + new TestNode( + "router", + (ctx, input) -> { + routerRuns.incrementAndGet(); + ctx.routes(ImmutableList.of(yes)); + return Flowable.empty(); + }); + Node ask = longRunningAsk("ask", "ask_1"); + CountingNode unreached = new CountingNode("unreached", "no_out"); + WorkflowSession session = + new WorkflowSession( + workflow( + new Edge(START, router), + new Edge(router, ask, yes), + new Edge(router, unreached, new Route.Tag("no")), + new Edge(ask, captureInput("reply")))); + session.advance(go()); + + ImmutableList events = session.turn(answer("ask_1", ImmutableMap.of("v", 1))); + + assertThat(routerRuns.get()).isEqualTo(1); + assertThat(unreached.runs()).isEqualTo(0); + assertThat(replyInput(events)).isEqualTo(ImmutableMap.of("v", 1)); + } + + @Test + public void resume_twoParallelInterruptsAnsweredOnSeparateTurns_joinsOnlyOnceBothAreAnswered() { + Node askA = longRunningAsk("ask_a", "a_1"); + Node askB = longRunningAsk("ask_b", "b_1"); + JoinNode join = JoinNode.builder().name("j").build(); + WorkflowSession session = + new WorkflowSession( + workflow( + new Edge(START, askA), + new Edge(START, askB), + new Edge(askA, join), + new Edge(askB, join))); + session.advance(go()); + + ImmutableList second = session.turn(answer("a_1", ImmutableMap.of("a", 1))); + ImmutableList third = session.turn(answer("b_1", ImmutableMap.of("b", 2))); + + assertThat(pathsOf(second)).doesNotContain("wf@1/j@1"); + assertThat(outputEventAt(third, "wf@1/j@1").output()) + .hasValue( + ImmutableMap.of("ask_a", ImmutableMap.of("a", 1), "ask_b", ImmutableMap.of("b", 2))); + } + + @Test + public void resumableWorkflow_run_checkpointsEachNodeAndEndsWithEndOfAgent() { + Node prep = emitter("prep", "prep_out"); + + ImmutableList events = new WorkflowSession(workflow(new Edge(START, prep))).turn(go()); + + ImmutableList> checkpoints = + events.stream() + .flatMap(event -> event.actions().agentState().stream()) + .collect(toImmutableList()); + assertThat(checkpoints) + .containsExactly( + ImmutableMap.of( + NodeState.NODES_KEY, ImmutableMap.of("prep", checkpoint(NodeStatus.RUNNING))), + ImmutableMap.of( + NodeState.NODES_KEY, ImmutableMap.of("prep", checkpoint(NodeStatus.COMPLETED)))) + .inOrder(); + Event last = Iterables.getLast(events); + assertThat(last.actions().endOfAgent()).isTrue(); + assertThat(last.author()).isEqualTo("wf"); + } + + @Test + public void workflowWithoutResumability_run_emitsNoCheckpointOrEndOfAgent() { + Node prep = emitter("prep", "prep_out"); + + ImmutableList events = + new WorkflowSession(workflow(new Edge(START, prep)), /* resumable= */ false).turn(go()); + + assertThat(events.stream().filter(event -> event.actions().agentState().isPresent())).isEmpty(); + assertThat(events.stream().filter(event -> event.actions().endOfAgent())).isEmpty(); + } + + @Test + public void resumableWorkflow_pausedRun_checkpointsTheWaitingNodeWithoutEndOfAgent() { + Node ask = longRunningAsk("ask", "ask_1"); + + ImmutableList events = new WorkflowSession(workflow(new Edge(START, ask))).turn(go()); + + Map lastCheckpoint = + Streams.findLast(events.stream().flatMap(event -> event.actions().agentState().stream())) + .orElseThrow(); + assertThat(lastCheckpoint) + .containsExactly( + NodeState.NODES_KEY, + ImmutableMap.of( + "ask", + ImmutableMap.of( + "status", NodeStatus.WAITING.code(), "interrupts", ImmutableList.of("ask_1")))); + assertThat(events.stream().filter(event -> event.actions().endOfAgent())).isEmpty(); + } + + @Test + public void resumableWorkflow_nodeWaitingForAnOutputItDidNotProduce_isCheckpointedAsWaiting() { + Node waiter = + new WaitingNode("waiter", (ctx, input) -> Flowable.empty(), /* rerunOnResume= */ false); + + ImmutableList events = new WorkflowSession(workflow(new Edge(START, waiter))).turn(go()); + + Map lastCheckpoint = + Streams.findLast(events.stream().flatMap(event -> event.actions().agentState().stream())) + .orElseThrow(); + assertThat(lastCheckpoint) + .containsExactly( + NodeState.NODES_KEY, ImmutableMap.of("waiter", checkpoint(NodeStatus.WAITING))); + } + + @Test + public void resumableWorkflow_canceledOnANodesCheckpoint_doesNotStartTheNode() { + CountingNode first = new CountingNode("first", "out"); + CountingNode second = new CountingNode("second", "out"); + WorkflowSession session = + new WorkflowSession(workflow(new Edge(START, first), new Edge(first, second))); + + ImmutableList events = session.turnUntil(go(), event -> marksRunning(event, "second")); + + assertThat(Iterables.getLast(events).actions().agentState()).isPresent(); + assertThat(first.runs()).isEqualTo(1); + assertThat(second.runs()).isEqualTo(0); + } + + @Test + public void resumableWorkflow_node_startsOnceItsCheckpointIsProcessed() { + CountingNode first = new CountingNode("first", "out"); + CountingNode second = new CountingNode("second", "out"); + AtomicInteger secondRunsAtItsCheckpoint = new AtomicInteger(-1); + Plugin asyncCheckpoints = + new BasePlugin("async_checkpoints") { + @Override + public Maybe onEventCallback(InvocationContext invocationContext, Event event) { + if (!marksRunning(event, "second")) { + return Maybe.empty(); + } + // Processed on another thread, so a node that did not wait could start meanwhile. + return Maybe.fromAction(() -> secondRunsAtItsCheckpoint.set(second.runs())) + .subscribeOn(Schedulers.single()); + } + }; + WorkflowSession session = + new WorkflowSession( + workflow(new Edge(START, first), new Edge(first, second)), + /* resumable= */ true, + asyncCheckpoints); + + ImmutableList events = session.turn(go()); + + assertThat(secondRunsAtItsCheckpoint.get()).isEqualTo(0); + assertThat(outputEventAt(events, "wf@1/second@1").output()).hasValue("out"); + } + + @Test + public void resume_finishedWorkflowInvocation_runsNothing() { + CountingNode prep = new CountingNode("prep", "prep_out"); + WorkflowSession session = new WorkflowSession(workflow(new Edge(START, prep))); + String invocationId = session.turn(go()).get(0).invocationId(); + + ImmutableList events = session.resume(invocationId); + + assertThat(events).isEmpty(); + assertThat(prep.runs()).isEqualTo(1); + } + + @Test + public void resume_textAnswerUnderResult_isReadAsJson() { + Node ask = longRunningAsk("ask", "ask_1"); + WorkflowSession session = + new WorkflowSession(workflow(new Edge(START, ask), new Edge(ask, captureInput("reply")))); + session.advance(go()); + + ImmutableList events = + session.turn(answer("ask_1", ImmutableMap.of("result", "{\"v\": 1}"))); + + assertThat(replyInput(events)).isEqualTo(ImmutableMap.of("v", 1)); + } + + @Test + public void resume_textAnswerUnderResultThatIsNotJson_staysText() { + Node ask = longRunningAsk("ask", "ask_1"); + WorkflowSession session = + new WorkflowSession(workflow(new Edge(START, ask), new Edge(ask, captureInput("reply")))); + session.advance(go()); + + ImmutableList events = + session.turn(answer("ask_1", ImmutableMap.of("result", "not json"))); + + assertThat(replyInput(events)).isEqualTo("not json"); + } + + @Test + public void resume_scalarAnswerUnderResult_isUnwrapped() { + Node ask = longRunningAsk("ask", "ask_1"); + WorkflowSession session = + new WorkflowSession(workflow(new Edge(START, ask), new Edge(ask, captureInput("reply")))); + session.advance(go()); + + ImmutableList events = session.turn(answer("ask_1", ImmutableMap.of("result", 5))); + + assertThat(replyInput(events)).isEqualTo(5); + } + + @Test + public void resume_nestedWorkflow_outputsItsTerminalNodesResult() { + Node step1 = emitter("step1", "X"); + Node ask = longRunningAsk("ask", "ask_1"); + Node done = emitter("done", "DONE"); + Workflow sub = + workflow("sub", new Edge(START, step1), new Edge(step1, ask), new Edge(ask, done)); + WorkflowSession session = + new WorkflowSession(workflow(new Edge(START, sub), new Edge(sub, captureInput("after")))); + session.advance(go()); + + ImmutableList events = session.turn(answer("ask_1", ImmutableMap.of("v", 1))); + + assertThat(outputEventAt(events, "wf@1/after@1").output()) + .hasValue(ImmutableMap.of("received", "DONE")); + } + + @Test + public void resume_valueEmittedBeforeTheInterrupt_isReplacedByTheAnswer() { + Node ask = new TestNode("ask", (ctx, input) -> Flowable.just("stale", interrupt("ask_1"))); + WorkflowSession session = + new WorkflowSession(workflow(new Edge(START, ask), new Edge(ask, captureInput("reply")))); + session.advance(go()); + + ImmutableList events = session.turn(answer("ask_1", ImmutableMap.of("v", 1))); + + assertThat(replyInput(events)).isEqualTo(ImmutableMap.of("v", 1)); + } + + @Test + public void resume_dynamicChildThatAsked_completesWithTheAnswer() { + Node dispatcher = + dispatcher( + "dispatch", + (ctx, input) -> + ctx.runNode(longRunningAsk("child", "ask_1"), null, runId("c")).toFlowable()); + WorkflowSession session = + new WorkflowSession( + workflow(new Edge(START, dispatcher), new Edge(dispatcher, captureInput("reply")))); + session.advance(go()); + + ImmutableList events = session.turn(answer("ask_1", ImmutableMap.of("v", 1))); + + assertThat(replyInput(events)).isEqualTo(ImmutableMap.of("v", 1)); + } + + @Test + public void resume_dynamicChildFinishedEarlier_replaysItsOutputWithoutRunning() { + CountingNode first = new CountingNode("first", "a_out"); + Node second = longRunningAsk("second", "ask_1"); + Node dispatcher = + dispatcher( + "dispatch", + (ctx, input) -> + ctx.runNode(first, null, runId("a")) + .flatMap( + a -> ctx.runNode(second, null, runId("b")).map(b -> ImmutableList.of(a, b))) + .toFlowable()); + WorkflowSession session = + new WorkflowSession( + workflow(new Edge(START, dispatcher), new Edge(dispatcher, captureInput("reply")))); + session.advance(go()); + + ImmutableList events = session.turn(answer("ask_1", ImmutableMap.of("v", 1))); + + assertThat(first.runs()).isEqualTo(1); + assertThat(replyInput(events)).isEqualTo(ImmutableList.of("a_out", ImmutableMap.of("v", 1))); + } + + @Test + public void resume_delegatedChildCompletedByTheAnswer_becomesTheCallersOutput() { + Node caller = + dispatcher( + "caller", + (ctx, input) -> + ctx.runNode(longRunningAsk("child", "ask_1"), null, USE_AS_OUTPUT) + .ignoreElement() + .toFlowable()); + WorkflowSession session = + new WorkflowSession( + workflow(new Edge(START, caller), new Edge(caller, captureInput("reply")))); + session.advance(go()); + + ImmutableList events = session.turn(answer("ask_1", ImmutableMap.of("v", 1))); + + assertThat(events.stream().filter(event -> event.errorCode().isPresent())).isEmpty(); + assertThat(outputEventAt(events, "wf@1/caller@1").output()).hasValue(ImmutableMap.of("v", 1)); + assertThat(replyInput(events)).isEqualTo(ImmutableMap.of("v", 1)); + } + + @Test + public void runNode_delegatingCallWhoseChildNoLongerRehydrates_letsTheCallerProduceItsOutput() { + Node chat = new MessageOutputNode("chat", "not json", AGE_SCHEMA); + Node ask = longRunningAsk("ask", "ask_1"); + // On resume, the replayed output of chat no longer parses, and the caller falls back. + Node caller = + dispatcher( + "caller", + (ctx, input) -> + ctx.runNode(chat, null, USE_AS_OUTPUT) + .ignoreElement() + .onErrorComplete() + .andThen(ctx.runNode(ask, null, runId("a"))) + .map(answer -> "fallback") + .toFlowable()); + WorkflowSession session = new WorkflowSession(workflow(new Edge(START, caller))); + session.advance(go()); + + ImmutableList events = session.turn(answer("ask_1", ImmutableMap.of("v", 1))); + + assertThat(outputEventAt(events, "wf@1/caller@1").output()).hasValue("fallback"); + } + + @Test + public void resume_rootWorkflowWhoseTerminalNodeReplayed_sendsNoOutputOfItsOwn() { + CountingNode done = new CountingNode("done", "DONE"); + Node sink = new TestNode("sink", (ctx, input) -> Flowable.empty()); + Node ask = longRunningAsk("ask", "ask_1"); + WorkflowSession session = + new WorkflowSession( + workflow(new Edge(START, done), new Edge(START, ask), new Edge(ask, sink))); + session.advance(go()); + + ImmutableList events = session.turn(answer("ask_1", ImmutableMap.of("v", 1))); + + assertThat(done.runs()).isEqualTo(1); + assertThat(outputEventAt(events, "wf@1/done@1").output()).hasValue("DONE"); + assertThat(events.stream().filter(event -> pathOf(event).equals("wf@1"))).isEmpty(); + } + + @Test + public void resume_callerRerunOverAReplayingChildWorkflow_isReplayedOnTheNextTurn() { + AtomicInteger callerRuns = new AtomicInteger(); + Workflow child = workflow("sub", new Edge(START, emitter("done", "DONE"))); + Node caller = + dispatcher( + "caller", + (ctx, input) -> { + callerRuns.incrementAndGet(); + return ctx.runNode(child, null, USE_AS_OUTPUT) + .ignoreElement() + .andThen(ctx.runNode(longRunningAsk("ask", "ask_1"), null, runId("a"))) + .ignoreElement() + .toFlowable(); + }); + Node other = longRunningAsk("other", "o_1"); + JoinNode join = JoinNode.builder().name("j").build(); + WorkflowSession session = + new WorkflowSession( + workflow( + new Edge(START, caller), + new Edge(START, other), + new Edge(caller, join), + new Edge(other, join))); + session.advance(go()); + session.advance(answer("ask_1", ImmutableMap.of("v", 1))); + + ImmutableList events = session.turn(answer("o_1", ImmutableMap.of("v", 2))); + + assertThat(callerRuns.get()).isEqualTo(2); + assertThat(outputEventAt(events, "wf@1/j@1").output()) + .hasValue(ImmutableMap.of("caller", "DONE", "other", ImmutableMap.of("v", 2))); + } + + @Test + public void resume_interruptRaisedByADescendant_reachesTheNodesResumeInputs() { + Node child = longRunningAsk("child", "ask_1"); + Node node = + dispatcher( + "node", + (ctx, input) -> + ctx.resumeInputs().isEmpty() + ? ctx.runNode(child, null, runId("c")).toFlowable() + : Flowable.just(ImmutableMap.of("resumed", ctx.resumeInputs()))); + WorkflowSession session = + new WorkflowSession(workflow(new Edge(START, node), new Edge(node, captureInput("reply")))); + session.advance(go()); + + ImmutableList events = session.turn(answer("ask_1", ImmutableMap.of("v", 1))); + + assertThat(replyInput(events)) + .isEqualTo(ImmutableMap.of("resumed", ImmutableMap.of("ask_1", ImmutableMap.of("v", 1)))); + } + + @Test + public void resume_concurrentDynamicChildren_replaysFinishedOnesAndCompletesTheAnsweredOne() { + Map runsByItem = new ConcurrentHashMap<>(); + CompletableSubject othersFinished1 = CompletableSubject.create(); + CompletableSubject othersFinished3 = CompletableSubject.create(); + Node worker = + new TestNode( + "worker", + (ctx, input) -> { + runsByItem.merge(input, 1, Integer::sum); + if (input.equals(2)) { + // Asks only once the other items have finished, so they have a result to replay. + return Completable.mergeArray(othersFinished1, othersFinished3) + .andThen(Flowable.just(interrupt("ask_2"))); + } + CompletableSubject finished = input.equals(1) ? othersFinished1 : othersFinished3; + return Flowable.just("done-" + input).doOnComplete(finished::onComplete); + }); + Node fanOut = + dispatcher( + "fan", + (ctx, input) -> + Flowable.fromIterable((List) input) + .concatMapEager( + item -> + ctx.runNode(worker, item, runId("item" + item)) + .subscribeOn(Schedulers.io()) + .toFlowable()) + .toList() + .toFlowable()); + Node items = emitter("items", ImmutableList.of(1, 2, 3)); + WorkflowSession session = + new WorkflowSession( + workflow( + new Edge(START, items), + new Edge(items, fanOut), + new Edge(fanOut, captureInput("reply")))); + session.advance(go()); + + ImmutableList events = session.turn(answer("ask_2", ImmutableMap.of("result", "two"))); + + assertThat(runsByItem).containsExactly(1, 1, 2, 1, 3, 1); + assertThat(replyInput(events)).isEqualTo(ImmutableList.of("done-1", "two", "done-3")); + } + + @Test + public void resume_rerunOnResumeNodeThatFinishedAfterItsAnswer_doesNotRunAgain() { + AtomicInteger writerRuns = new AtomicInteger(); + Node writer = + dispatcher( + "writer", + (ctx, input) -> { + writerRuns.incrementAndGet(); + if (ctx.resumeInputs().isEmpty()) { + return Flowable.just(interrupt("ask_w")); + } + ctx.state().put("writer", "done"); + return Flowable.empty(); + }); + Node other = longRunningAsk("other", "ask_o"); + JoinNode join = JoinNode.builder().name("j").build(); + WorkflowSession session = + new WorkflowSession( + workflow( + new Edge(START, writer), + new Edge(START, other), + new Edge(writer, join), + new Edge(other, join))); + session.advance(go()); + session.advance(answer("ask_w", ImmutableMap.of("v", 1))); + + ImmutableList events = session.turn(answer("ask_o", ImmutableMap.of("v", 2))); + + assertThat(writerRuns.get()).isEqualTo(2); + assertThat(pathsOf(events)).contains("wf@1/j@1"); + } + + @Test + public void resume_rerunOnResumeNodeThatFinishedWithoutAnOutput_doesNotRunAgain() { + AtomicInteger writerRuns = new AtomicInteger(); + Node writer = + dispatcher( + "writer", + (ctx, input) -> { + writerRuns.incrementAndGet(); + ctx.state().put("writer", "done"); + return Flowable.empty(); + }); + Node ask = longRunningAsk("ask", "ask_1"); + WorkflowSession session = + new WorkflowSession(workflow(new Edge(START, writer), new Edge(START, ask))); + session.advance(go()); + + session.advance(answer("ask_1", ImmutableMap.of("v", 1))); + + assertThat(writerRuns.get()).isEqualTo(1); + } + + @Test + public void resume_nodeAfterStartThatRunsAgain_getsTheOriginalUserMessage() { + List<@Nullable Object> inputs = Collections.synchronizedList(new ArrayList<>()); + Node reader = + dispatcher( + "reader", + (ctx, input) -> { + inputs.add(input); + return ctx.resumeInputs().isEmpty() + ? Flowable.just(interrupt("ask_r")) + : Flowable.just("done"); + }); + WorkflowSession session = new WorkflowSession(workflow(new Edge(START, reader))); + session.advance(go()); + + session.advance(answer("ask_r", ImmutableMap.of("v", 1))); + + assertThat(inputs).containsExactly(go(), go()); + } + + @Test + public void resume_afterAnEarlierInvocation_nodeAfterStartGetsItsInvocationsMessage() { + Content first = Content.builder().role("user").parts(Part.fromText("first")).build(); + Content second = Content.builder().role("user").parts(Part.fromText("second")).build(); + List<@Nullable Object> inputs = Collections.synchronizedList(new ArrayList<>()); + Node reader = + dispatcher( + "reader", + (ctx, input) -> { + inputs.add(input); + return Objects.equals(input, second) && ctx.resumeInputs().isEmpty() + ? Flowable.just(interrupt("ask_r")) + : Flowable.just("done"); + }); + WorkflowSession session = new WorkflowSession(workflow(new Edge(START, reader))); + session.advance(first); + session.advance(second); + + session.advance(answer("ask_r", ImmutableMap.of("v", 1))); + + assertThat(inputs).containsExactly(first, second, second).inOrder(); + } + + @Test + public void resume_nodeThatRunsAgain_seesTheOriginalUserMessageAsItsUserContent() { + List<@Nullable Content> userContents = Collections.synchronizedList(new ArrayList<>()); + Node reader = + dispatcher( + "reader", + (ctx, input) -> { + userContents.add(ctx.userContent().orElse(null)); + return ctx.resumeInputs().isEmpty() + ? Flowable.just(interrupt("ask_r")) + : Flowable.just("done"); + }); + WorkflowSession session = new WorkflowSession(workflow(new Edge(START, reader))); + session.advance(go()); + + session.advance(answer("ask_r", ImmutableMap.of("v", 1))); + + assertThat(userContents).containsExactly(go(), go()); + } + + @Test + public void resume_standaloneRootNodeAnswered_completesWithTheAnswer() { + WorkflowSession session = new WorkflowSession(longRunningAsk("ask", "ask_1")); + session.advance(go()); + + ImmutableList events = session.turn(answer("ask_1", ImmutableMap.of("v", 1))); + + assertThat(outputEventAt(events, "ask@1").output()).hasValue(ImmutableMap.of("v", 1)); + assertThat(interruptEvents(events)).isEmpty(); + } + + @Test + public void resume_llmAgentNodeAnsweredLongRunningCall_callsTheModelAgain() { + TestLlm model = + TestLlm.create( + TestUtils.createFunctionCallLlmResponse("call_1", "approve", ImmutableMap.of()), + TestUtils.createTextLlmResponse("approved")); + LlmAgent agent = + LlmAgent.builder().name("agent").model(model).tools(new PendingTool("approve")).build(); + WorkflowSession session = new WorkflowSession(workflow(new Edge(START, agent))); + session.advance(go()); + + ImmutableList events = + session.turn(answer("call_1", ImmutableMap.of("ok", true), "approve")); + + ImmutableList requests = model.getRequests(); + assertThat(requests).hasSize(2); + assertThat( + requests.get(1).contents().stream() + .flatMap(content -> content.parts().orElse(ImmutableList.of()).stream()) + .flatMap(part -> part.functionResponse().stream()) + .map(response -> response.id().orElse(""))) + .contains("call_1"); + assertThat(events.stream().map(Event::stringifyContent)).contains("approved"); + } + + @Test + public void resume_rolelessAnswerWithoutResumability_reachesTheModel() { + TestLlm model = + TestLlm.create( + TestUtils.createFunctionCallLlmResponse("call_1", "approve", ImmutableMap.of()), + TestUtils.createTextLlmResponse("approved")); + LlmAgent agent = + LlmAgent.builder().name("agent").model(model).tools(new PendingTool("approve")).build(); + WorkflowSession session = + new WorkflowSession(workflow(new Edge(START, agent)), /* resumable= */ false); + session.advance(go()); + Content withRole = answer("call_1", ImmutableMap.of("ok", true), "approve"); + + ImmutableList events = + session.turn(Content.builder().parts(withRole.parts().orElseThrow()).build()); + + ImmutableList requests = model.getRequests(); + assertThat(requests).hasSize(2); + assertThat( + requests.get(1).contents().stream() + .flatMap(content -> content.parts().orElse(ImmutableList.of()).stream()) + .flatMap(part -> part.functionResponse().stream()) + .map(response -> response.id().orElse(""))) + .contains("call_1"); + assertThat(events.stream().map(Event::stringifyContent)).contains("approved"); + } + + @Test + public void resume_messageOutputNode_isRehydratedToItsText() { + Node chat = new MessageOutputNode("chat", "agent said hi", /* schema= */ null); + Node ask = longRunningAsk("ask", "ask_1"); + JoinNode join = JoinNode.builder().name("j").build(); + WorkflowSession session = + new WorkflowSession( + workflow( + new Edge(START, chat), + new Edge(START, ask), + new Edge(chat, join), + new Edge(ask, join))); + session.advance(go()); + + ImmutableList events = session.turn(answer("ask_1", ImmutableMap.of("v", 1))); + + assertThat(outputEventAt(events, "wf@1/j@1").output()) + .hasValue(ImmutableMap.of("chat", "agent said hi", "ask", ImmutableMap.of("v", 1))); + } + + @Test + public void resume_messageOutputNodeWithAnOutputSchema_isRehydratedToTheParsedValue() { + Schema schema = + Schema.builder() + .type(Type.Known.OBJECT) + .properties(ImmutableMap.of("city", Schema.builder().type(Type.Known.STRING).build())) + .build(); + Node chat = new MessageOutputNode("chat", "{\"city\": \"Paris\"}", schema); + Node ask = longRunningAsk("ask", "ask_1"); + JoinNode join = JoinNode.builder().name("j").build(); + WorkflowSession session = + new WorkflowSession( + workflow( + new Edge(START, chat), + new Edge(START, ask), + new Edge(chat, join), + new Edge(ask, join))); + session.advance(go()); + + ImmutableList events = session.turn(answer("ask_1", ImmutableMap.of("v", 1))); + + assertThat(outputEventAt(events, "wf@1/j@1").output()) + .hasValue( + ImmutableMap.of( + "chat", ImmutableMap.of("city", "Paris"), "ask", ImmutableMap.of("v", 1))); + } + + @Test + public void resume_nodeThatFailedInAnEarlierTurn_runsAgain() { + AtomicInteger flakyRuns = new AtomicInteger(); + Node flaky = + new TestNode( + "flaky", + (ctx, input) -> + flakyRuns.incrementAndGet() == 1 + ? Flowable.error(new IllegalStateException("first run fails")) + : Flowable.just("recovered")); + Node ask = longRunningAsk("ask", "ask_1"); + JoinNode join = JoinNode.builder().name("j").build(); + WorkflowSession session = + new WorkflowSession( + workflow( + new Edge(START, ask), + new Edge(START, flaky), + new Edge(ask, join), + new Edge(flaky, join))); + FailedTurn first = session.failingTurn(go()); + + ImmutableList events = session.turn(answer("ask_1", ImmutableMap.of("v", 1))); + + assertThat(first.error()).isInstanceOf(IllegalStateException.class); + assertThat(outputEventAt(events, "wf@1/j@1").output()) + .hasValue(ImmutableMap.of("ask", ImmutableMap.of("v", 1), "flaky", "recovered")); + } + + @Test + public void resume_nodeCompletingATransfer_checkpointsItAsRunning() { + Node node = + new TestNode( + "node", + (ctx, input) -> + Flowable.just( + newEvent() + .content(callContent("approve", "t_1")) + .longRunningToolIds(ImmutableSet.of("t_1")) + .actions(EventActions.builder().transferToAgent("other").build()) + .build())); + WorkflowSession session = new WorkflowSession(workflow(new Edge(START, node))); + session.advance(go()); + + ImmutableList events = + session.turn(answer("t_1", ImmutableMap.of("result", "ok"), "approve")); + + assertThat(events.stream().flatMap(event -> event.actions().agentState().stream())) + .contains( + ImmutableMap.of( + NodeState.NODES_KEY, ImmutableMap.of("node", checkpoint(NodeStatus.RUNNING)))); + } + + @Test + public void resume_fannedOutNodeReplayed_keepsItsSubBranchOnTheEventItEmitsAgain() { + CountingNode done = new CountingNode("done", "done_out"); + Node ask = longRunningAsk("ask", "ask_1"); + JoinNode join = JoinNode.builder().name("j").build(); + WorkflowSession session = + new WorkflowSession( + workflow( + new Edge(START, done), + new Edge(START, ask), + new Edge(done, join), + new Edge(ask, join))); + session.advance(go()); + + ImmutableList events = session.turn(answer("ask_1", ImmutableMap.of("v", 1))); + + assertThat(done.runs()).isEqualTo(1); + assertThat(outputEventAt(events, "wf@1/done@1").branch()).hasValue("done@1"); + } + + @Test + public void resume_nodeEnteredAgainAfterTheResume_takesTheNextRunId() { + AtomicInteger askRuns = new AtomicInteger(); + Node ask = + new TestNode( + "ask", (ctx, input) -> Flowable.just(interrupt("ask_" + askRuns.incrementAndGet()))); + AtomicInteger againRuns = new AtomicInteger(); + Route yes = new Route.Tag("yes"); + Node again = + new TestNode( + "again", + (ctx, input) -> { + if (againRuns.incrementAndGet() == 1) { + ctx.routes(ImmutableList.of(yes)); + } + return Flowable.empty(); + }); + WorkflowSession session = + new WorkflowSession( + workflow(new Edge(START, ask), new Edge(ask, again), new Edge(again, ask, yes))); + session.advance(go()); + + ImmutableList events = session.turn(answer("ask_1", ImmutableMap.of("v", 1))); + + assertThat(outputEventAt(events, "wf@1/ask@1").output()).hasValue(ImmutableMap.of("v", 1)); + Event request = interruptEvents(events).stream().collect(onlyElement()); + assertThat(pathOf(request)).isEqualTo("wf@1/ask@2"); + assertThat(request.functionCalls().get(0).id()).hasValue("ask_2"); + } + + @Test + public void runNode_failedChildWorkflowRunAgain_rerunsTheSiblingThatItsFailureCanceled() { + AtomicInteger aRuns = new AtomicInteger(); + AtomicInteger bRuns = new AtomicInteger(); + CompletableSubject bStarted = CompletableSubject.create(); + Node a = + new TestNode( + "a", + (ctx, input) -> + aRuns.incrementAndGet() == 1 + ? bStarted.andThen(Flowable.error(new IllegalStateException("a fails"))) + : Flowable.empty()); + // On its first run, b reports progress and then works until the failure of a cancels it. + Node b = + new TestNode( + "b", + (ctx, input) -> + bRuns.incrementAndGet() == 1 + ? Flowable.just(newEvent().content(modelText("working")).build()) + .concatWith(Completable.fromAction(bStarted::onComplete)) + .concatWith(Flowable.never()) + : Flowable.just("b")); + Workflow inner = workflow("inner", new Edge(START, a), new Edge(START, b)); + Node caller = + dispatcher( + "caller", + (ctx, input) -> + ctx.runNode(inner, null, runId("w")) + .onErrorResumeNext( + error -> + error instanceof DynamicNodeFailedException + ? ctx.runNode(inner, null, runId("w")) + : Maybe.error(error)) + .toFlowable()); + + ImmutableList events = new WorkflowSession(workflow(new Edge(START, caller))).turn(go()); + + assertThat(aRuns.get()).isEqualTo(2); + assertThat(bRuns.get()).isEqualTo(2); + assertThat(outputEventAt(events, "wf@1/caller@1").output()).hasValue("b"); + } + + @Test + public void runNode_failedChildWorkflowRunAgain_doesNotRerunANodeThatFinishedWithoutAnOutput() { + AtomicInteger cRuns = new AtomicInteger(); + AtomicInteger aRuns = new AtomicInteger(); + // c emits no output; only being marked finished keeps it from counting as canceled mid-run. + Node c = + new TestNode( + "c", + (ctx, input) -> { + cRuns.incrementAndGet(); + return Flowable.just(newEvent().content(modelText("working")).build()); + }); + Node a = + new TestNode( + "a", + (ctx, input) -> + aRuns.incrementAndGet() == 1 + ? Flowable.error(new IllegalStateException("a fails")) + : Flowable.just("a")); + Workflow inner = workflow("inner", new Edge(START, c), new Edge(c, a)); + Node caller = + dispatcher( + "caller", + (ctx, input) -> + ctx.runNode(inner, null, runId("w")) + .onErrorResumeNext( + error -> + error instanceof DynamicNodeFailedException + ? ctx.runNode(inner, null, runId("w")) + : Maybe.error(error)) + .toFlowable()); + + ImmutableList events = new WorkflowSession(workflow(new Edge(START, caller))).turn(go()); + + assertThat(cRuns.get()).isEqualTo(1); + assertThat(aRuns.get()).isEqualTo(2); + assertThat(outputEventAt(events, "wf@1/caller@1").output()).hasValue("a"); + } + + @Test + public void runNode_failedCallerRunAgain_rerunsTheChildItsFailureCanceled() { + AtomicInteger aRuns = new AtomicInteger(); + AtomicInteger bRuns = new AtomicInteger(); + CompletableSubject bStarted = CompletableSubject.create(); + Node a = + new TestNode( + "a", + (ctx, input) -> + aRuns.incrementAndGet() == 1 + ? bStarted.andThen(Flowable.error(new IllegalStateException("a fails"))) + : Flowable.just("a")); + // b reruns on resume, so its progress event alone would read as a finished run. + Node b = + dispatcher( + "b", + (ctx, input) -> + bRuns.incrementAndGet() == 1 + ? Flowable.just(newEvent().content(modelText("working")).build()) + .concatWith(Completable.fromAction(bStarted::onComplete)) + .concatWith(Flowable.never()) + : Flowable.just("b")); + Node inner = + dispatcher( + "inner", + (ctx, input) -> + Maybe.zip( + ctx.runNode(b, null, runId("b")), + ctx.runNode(a, null, runId("a")), + (bOutput, aOutput) -> bOutput) + .toFlowable()); + Node caller = + dispatcher( + "caller", + (ctx, input) -> + ctx.runNode(inner, null, runId("w")) + .onErrorResumeNext( + error -> + error instanceof DynamicNodeFailedException + ? ctx.runNode(inner, null, runId("w")) + : Maybe.error(error)) + .toFlowable()); + + ImmutableList events = new WorkflowSession(workflow(new Edge(START, caller))).turn(go()); + + assertThat(aRuns.get()).isEqualTo(2); + assertThat(bRuns.get()).isEqualTo(2); + assertThat(outputEventAt(events, "wf@1/caller@1").output()).hasValue("b"); + } + + private static ImmutableList interruptEvents(List events) { + return events.stream() + .filter(event -> !event.longRunningToolIds().orElse(ImmutableSet.of()).isEmpty()) + .collect(toImmutableList()); + } + + /** Returns whether {@code event} is a checkpoint that marks {@code nodeName} as running. */ + private static boolean marksRunning(Event event, String nodeName) { + return event + .actions() + .agentState() + .map(state -> (Map) state.get(NodeState.NODES_KEY)) + .filter(nodes -> checkpoint(NodeStatus.RUNNING).equals(nodes.get(nodeName))) + .isPresent(); + } + + private static ImmutableMap checkpoint(NodeStatus status) { + return ImmutableMap.of("status", status.code(), "interrupts", ImmutableList.of()); + } + + private static RunNodeConfig runId(String runId) { + return RunNodeConfig.builder().runId(runId).build(); + } + + /** Returns a node that waits on the long-running call {@code interruptId}, as a tool does. */ + private static Node longRunningAsk(String name, String interruptId) { + return new TestNode(name, (ctx, input) -> Flowable.just(interrupt(interruptId))); + } + + /** Returns the event in which a node waits on the long-running call {@code interruptId}. */ + private static Event interrupt(String interruptId) { + return newEvent() + .content(callContent(LONG_RUNNING_CALL, interruptId)) + .longRunningToolIds(ImmutableSet.of(interruptId)) + .build(); + } + + /** A node that emits {@code output}, if any, and counts its runs. */ + private static final class CountingNode implements Node { + private final String name; + private final @Nullable Object output; + private final AtomicInteger runs = new AtomicInteger(); + + CountingNode(String name, @Nullable Object output) { + this.name = name; + this.output = output; + } + + int runs() { + return runs.get(); + } + + @Override + public String name() { + return name; + } + + @Override + public Flowable runNode(Context context, @Nullable Object nodeInput) { + runs.incrementAndGet(); + return output == null ? Flowable.empty() : Flowable.just(output); + } + } + + /** A long-running tool that answers nothing, so the call stays pending until the user answers. */ + private static final class PendingTool extends BaseTool { + PendingTool(String name) { + super(name, "Asks the user to approve.", /* isLongRunning= */ true); + } + + @Override + public Optional declaration() { + return Optional.of( + FunctionDeclaration.builder().name(name()).description(description()).build()); + } + + @Override + public Single> runAsync(Map args, ToolContext toolContext) { + return Single.just(ImmutableMap.of()); + } + } +} diff --git a/core/src/test/java/com/google/adk/workflow/WorkflowResumeTesting.java b/core/src/test/java/com/google/adk/workflow/WorkflowResumeTesting.java index 234b49ed8..f6fd857e0 100644 --- a/core/src/test/java/com/google/adk/workflow/WorkflowResumeTesting.java +++ b/core/src/test/java/com/google/adk/workflow/WorkflowResumeTesting.java @@ -16,8 +16,20 @@ package com.google.adk.workflow; +import static com.google.common.base.Preconditions.checkState; +import static com.google.common.collect.ImmutableList.toImmutableList; +import static com.google.common.collect.MoreCollectors.onlyElement; +import static java.util.concurrent.TimeUnit.SECONDS; + import com.google.adk.agents.Context; +import com.google.adk.agents.RunConfig; +import com.google.adk.apps.App; +import com.google.adk.apps.ResumabilityConfig; import com.google.adk.events.Event; +import com.google.adk.plugins.Plugin; +import com.google.adk.runner.Runner; +import com.google.adk.testing.TestNode; +import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.genai.types.Content; import com.google.genai.types.FunctionCall; @@ -26,13 +38,22 @@ import com.google.genai.types.Schema; import com.google.genai.types.Type; import io.reactivex.rxjava3.core.Flowable; +import io.reactivex.rxjava3.functions.Predicate; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.concurrent.atomic.AtomicReference; import java.util.function.BiFunction; import org.jspecify.annotations.Nullable; /** Helpers for the tests that pause a workflow and resume it. */ final class WorkflowResumeTesting { + + private static final String APP = "resume_app"; + private static final String USER = "user"; + static final Start START = Start.INSTANCE; // The name of the long-running call that the nodes of these tests wait on. static final String LONG_RUNNING_CALL = "approve"; @@ -46,6 +67,30 @@ final class WorkflowResumeTesting { private WorkflowResumeTesting() {} + /** Returns the input that the {@code reply} node of a workflow named {@code wf} received. */ + static @Nullable Object replyInput(List events) { + Object output = outputEventAt(events, "wf@1/reply@1").output().orElseThrow(); + return ((Map) output).get("received"); + } + + static Event outputEventAt(List events, String nodePath) { + return events.stream() + .filter(event -> pathOf(event).equals(nodePath) && event.output().isPresent()) + .collect(onlyElement()); + } + + static ImmutableList pathsOf(List events) { + return events.stream().map(WorkflowResumeTesting::pathOf).collect(toImmutableList()); + } + + static String pathOf(Event event) { + return event.nodeInfo().map(NodeInfo::path).orElse(""); + } + + static Content go() { + return Content.builder().role("user").parts(Part.fromText("go")).build(); + } + static Content answer(String interruptId, Map response) { return answer(interruptId, response, LONG_RUNNING_CALL); } @@ -82,6 +127,94 @@ static Event.Builder newEvent() { return Event.builder().id(Event.generateEventId()).author(""); } + /** Returns a node that outputs the input it received under the key {@code received}. */ + static Node captureInput(String name) { + return new TestNode( + name, (ctx, input) -> Flowable.just(Collections.singletonMap("received", input))); + } + + /** A runner over an app rooted on a node, driven a turn at a time on one session. */ + static final class WorkflowSession { + private final Runner runner; + private final String sessionId; + + WorkflowSession(Node root) { + this(root, /* resumable= */ true); + } + + WorkflowSession(Node root, boolean resumable, Plugin... plugins) { + App.Builder app = App.builder().name(APP).rootNode(root).plugins(plugins); + if (resumable) { + app.resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()); + } + this.runner = Runner.builder().app(app.build()).build(); + this.sessionId = runner.sessionService().createSession(APP, USER).blockingGet().id(); + } + + ImmutableList turn(Content message) { + return ImmutableList.copyOf( + runner.runAsync(USER, sessionId, message).timeout(30, SECONDS).toList().blockingGet()); + } + + ImmutableList resume(String invocationId) { + return ImmutableList.copyOf( + runner + .runAsync( + USER, + sessionId, + invocationId, + /* newMessage= */ null, + RunConfig.builder().build(), + /* stateDelta= */ null) + .timeout(30, SECONDS) + .toList() + .blockingGet()); + } + + /** Runs a turn until an event matches {@code last}, which cancels the rest of the turn. */ + ImmutableList turnUntil(Content message, Predicate last) { + return ImmutableList.copyOf( + runner + .runAsync(USER, sessionId, message) + .takeUntil(last) + .timeout(30, SECONDS) + .toList() + .blockingGet()); + } + + /** Runs a turn that must fail, returning the events it streamed and its failure. */ + FailedTurn failingTurn(Content message) { + List events = new ArrayList<>(); + AtomicReference error = new AtomicReference<>(); + runner + .runAsync(USER, sessionId, message) + .timeout(30, SECONDS) + .blockingSubscribe(events::add, error::set); + checkState(error.get() != null, "expected the turn to fail"); + return new FailedTurn(ImmutableList.copyOf(events), error.get()); + } + + /** Runs a turn whose events matter only for the turns after it. */ + void advance(Content message) { + checkState(!turn(message).isEmpty(), "a turn produced no events to resume from"); + } + + ImmutableList invocationIds() { + return runner + .sessionService() + .getSession(APP, USER, sessionId, Optional.empty()) + .blockingGet() + .immutableEvents() + .stream() + .map(Event::invocationId) + .distinct() + .collect(toImmutableList()); + } + } + + /** The events a failed turn streamed before it failed, and its failure. */ + record FailedTurn(ImmutableList events, Throwable error) {} + /** A node that waits for an output and runs its {@code body}. */ record WaitingNode( String name, BiFunction> body, boolean rerunOnResume)