Skip to content

feat: resume a paused workflow from its session events #1688

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
copybara-service merged 1 commit into from
Oct 23, 2026
+1,959 −119
Merged

feat: resume a paused workflow from its session events #1688

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 13 additions & 0 deletions core/src/main/java/com/google/adk/agents/Context.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -321,6 +322,18 @@ public final void routes(List<? extends Route> routes) {
requireNodeState().selectRoutes(routes);
}

/**
* Returns the user's answers to the long-running calls that this workflow node or its descendants
* made in an earlier turn, keyed by interrupt ID. A node that reruns on resume reads these
* answers. A null answer has a null value in the map.
*
* @throws IllegalStateException if this is not the context of a workflow node.
*/
@Experimental
public final Map<String, @Nullable Object> resumeInputs() {
return requireNodeState().resumeInputs();
}

/**
* Runs {@code node} as a child of this workflow node without an input, as {@link #runNode(Node,
* Object, RunNodeConfig)} does with the default {@link RunNodeConfig}.
Expand Down
11 changes: 11 additions & 0 deletions core/src/main/java/com/google/adk/agents/LlmAgent.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -1082,6 +1083,16 @@ public Optional<Schema> outputSchema() {
return outputSchema;
}

/**
* Returns {@code true}, so that a workflow resuming this agent as a node runs it again and its
* model sees the answers to its long-running calls.
*/
@Experimental
@Override
public boolean rerunOnResume() {
return true;
}

public Optional<Executor> executor() {
return executor;
}
Expand Down
158 changes: 89 additions & 69 deletions core/src/main/java/com/google/adk/runner/Runner.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -111,8 +111,6 @@ public class Runner {
private final UuidProvider uuidProvider;
private final ConcurrentMap<String, Completable> activeSessionCompletables =
new MapMaker().weakValues().makeMap();
private static final String NODE_ROOT_RESUME_UNSUPPORTED =
"Resuming an invocation of a workflow node root is not supported.";

/** Builder for {@link Runner}. */
public static class Builder {
Expand Down Expand Up @@ -668,7 +666,8 @@ public Flowable<Event> runAsync(String userId, String sessionId, Content newMess
* Runs the agent, resuming an existing invocation instead of starting a new one. The invocation
* is resolved from {@code invocationId}, or from a function response carried by {@code
* newMessage}. Agent checkpoints are rehydrated from history and an invocation whose active agent
* already finished resolves to a no-op.
* already finished resolves to a no-op, while a workflow node root rebuilds its progress from the
* invocation's events even when the app is not resumable.
*
* @param userId the user id of the session.
* @param sessionId the session id.
Expand All @@ -679,7 +678,7 @@ public Flowable<Event> runAsync(String userId, String sessionId, Content newMess
* @param stateDelta optional state updates to merge into the session for this run.
* @return the events generated while resuming, or an empty stream when there is nothing to
* resume.
* @throws IllegalStateException if the app is not resumable.
* @throws IllegalStateException if the app is not resumable and its root is an agent.
* @throws IllegalArgumentException if the invocation cannot be resolved.
*/
@Experimental
Expand All @@ -691,7 +690,7 @@ public Flowable<Event> runAsync(
RunConfig runConfig,
@Nullable Map<String, Object> stateDelta) {
checkState(
isResumable(),
isResumable() || node != null,
"Resuming an invocation requires an App configured with a resumable ResumabilityConfig.");
Flowable<Event> result =
Flowable.defer(
Expand Down Expand Up @@ -782,26 +781,25 @@ protected Flowable<Event> runAsyncImpl(
Preconditions.checkNotNull(session, "session cannot be null");
Preconditions.checkNotNull(newMessage, "newMessage cannot be null");
Preconditions.checkNotNull(runConfig, "runConfig cannot be null");
if (isResumable()) {
// A node root rebuilds its progress from its invocation's events, so it resumes on any app.
if (isResumable() || node != null) {
return runResumableFromSession(
session, /* providedInvocationId= */ null, newMessage, runConfig, stateDelta);
}
return runNewInvocation(session, newMessage, runConfig, stateDelta);
}

/** Starts a brand-new invocation for {@code newMessage} (the default, non-resume flow). */
/**
* Starts an invocation for {@code newMessage}; a node root joins the one a plugin's answers
* resolve to.
*/
private Flowable<Event> runNewInvocation(
Session session,
Content newMessage,
RunConfig runConfig,
@Nullable Map<String, Object> stateDelta) {
return Flowable.defer(
() -> {
// Rejected before any callback or append; defer turns the throw into an error signal.
if (node != null) {
checkNoTextWithFunctionResponses(newMessage);
checkNoAnswerToAnEarlierInvocation(session, newMessage);
}
Context capturedContext = Context.current();
BaseAgent rootAgent = this.agent;
String invocationId = InvocationContext.newInvocationContextId(this.uuidProvider);
Expand All @@ -825,18 +823,21 @@ private Flowable<Event> runNewInvocation(
.onUserMessageCallback(initialContext, newMessage)
.compose(Tracing.<Content>withContext(capturedContext))
.defaultIfEmpty(newMessage)
.flatMap(
content ->
createUserMessageEvent(
session,
content,
initialContext,
runConfig.saveInputBlobsAsArtifacts(),
stateDelta))
.flatMapPublisher(
userEvent ->
runAgentForUserEvent(initialContext, session, userEvent, rootAgent)
.compose(Tracing.<Event>withContext(capturedContext)))
content -> {
InvocationContext context =
resolveNodeRootContext(session, initialContext, content);
return createUserMessageEvent(
session,
content,
context,
runConfig.saveInputBlobsAsArtifacts(),
stateDelta)
.flatMapPublisher(
userEvent ->
runAgentForUserEvent(context, session, userEvent, rootAgent)
.compose(Tracing.<Event>withContext(capturedContext)));
})
.doOnError(throwable -> runOnRunError(initialContext, throwable));
})
.doOnError(
Expand Down Expand Up @@ -893,7 +894,11 @@ private Flowable<Event> runAgentWithUpdatedSession(
initialContext.toBuilder()
.session(updatedSession)
.agent(agentToRun)
.userContent(event.content().orElseGet(Content::fromParts))
.userContent(
userContentFor(
updatedSession,
initialContext.invocationId(),
event.content().orElseGet(Content::fromParts)))
.branch(routedAgentParentBranch(updatedSession, agentToRun, rootAgent))
.build();

Expand Down Expand Up @@ -1006,7 +1011,7 @@ private void runOnRunError(InvocationContext context, Throwable throwable) {
/**
* Resumes an existing invocation when one resolves from {@code providedInvocationId} or a
* function response in {@code rawMessage}; otherwise starts a new invocation. Requires
* resumability.
* resumability or a workflow node root.
*/
private Flowable<Event> runResumableFromSession(
Session session,
Expand Down Expand Up @@ -1043,11 +1048,6 @@ private Flowable<Event> runResumableFromSession(
return Flowable.<Event>error(
new IllegalArgumentException("Session has no events to resume."));
}
if (node != null) {
// Without replay, resuming would run the node again and repeat its side effects.
return Flowable.<Event>error(
new UnsupportedOperationException(NODE_ROOT_RESUME_UNSUPPORTED));
}
return resumeCore(session, resolvedInvocationId, newMessage, runConfig, stateDelta);
});
}
Expand All @@ -1067,50 +1067,52 @@ private Flowable<Event> resumeCore(
() -> {
Context capturedContext = Context.current();

// Not new input: Python drops it and its state delta rather than adding a second prompt.
// Not new input: Python drops newMessage, and outside the node path stateDelta too.
boolean supersededByOriginalMessage =
newMessage != null
&& !carriesFunctionResponse(newMessage)
&& findOriginalUserMessage(session, resolvedInvocationId).isPresent();
Content message = supersededByOriginalMessage ? null : newMessage;
Map<String, Object> delta =
supersededByOriginalMessage && node == null ? null : stateDelta;

// Merged eagerly so the on-user-message callback sees it.
if (!supersededByOriginalMessage && stateDelta != null && !stateDelta.isEmpty()) {
stateDelta.forEach((key, value) -> session.state().put(key, value));
if (delta != null && !delta.isEmpty()) {
delta.forEach((key, value) -> session.state().put(key, value));
}

// For the pre-run plugin callbacks; the agent is resolved later, in runResumedAgent.
InvocationContext initialContext =
newInvocationContextBuilder(session)
.invocationId(resolvedInvocationId)
.runConfig(runConfig)
.userContent(
newMessage != null
? newMessage
: originalUserMessage(session, resolvedInvocationId))
.userContent(userContentFor(session, resolvedInvocationId, newMessage))
.build();

Flowable<Event> events;
if (supersededByOriginalMessage) {
events =
runResumedAgent(session, resolvedInvocationId, /* userContent= */ null, runConfig);
} else if (newMessage != null) {
if (message != null) {
// Inherit the branch of the call it answers, so routing and rehydration see it.
String branch =
matchingFunctionCallEvent(session, newMessage).flatMap(Event::branch).orElse(null);
matchingFunctionCallEvent(session, message).flatMap(Event::branch).orElse(null);
events =
this.pluginManager
.onUserMessageCallback(initialContext, newMessage)
.onUserMessageCallback(initialContext, message)
.compose(Tracing.<Content>withContext(capturedContext))
.defaultIfEmpty(newMessage)
.defaultIfEmpty(message)
.flatMap(
content ->
appendNewMessageToSession(
session,
content,
initialContext,
runConfig.saveInputBlobsAsArtifacts(),
stateDelta,
branch))
content -> {
// A plugin's message must not mix answers and text either.
if (node != null) {
checkNoTextWithFunctionResponses(content);
}
return appendNewMessageToSession(
session,
content,
initialContext,
runConfig.saveInputBlobsAsArtifacts(),
delta,
branch);
})
// Persist first: an unpersisted response leaves the call looking unanswered.
.flatMap(userEvent -> this.sessionService.appendEvent(session, userEvent))
.flatMapPublisher(
Expand All @@ -1120,7 +1122,7 @@ private Flowable<Event> resumeCore(
resolvedInvocationId,
userEvent.content().orElse(null),
runConfig));
} else if (stateDelta != null && !stateDelta.isEmpty()) {
} else if (delta != null && !delta.isEmpty()) {
// No message to carry the delta, so persist a content-less event, as Python does.
Event stateDeltaEvent =
Event.builder()
Expand All @@ -1129,9 +1131,7 @@ private Flowable<Event> resumeCore(
.invocationId(resolvedInvocationId)
.author(Role.USER)
.actions(
EventActions.builder()
.stateDelta(new ConcurrentHashMap<>(stateDelta))
.build())
EventActions.builder().stateDelta(new ConcurrentHashMap<>(delta)).build())
.build();
events =
this.sessionService
Expand Down Expand Up @@ -1168,10 +1168,7 @@ private Flowable<Event> runResumedAgent(
newInvocationContextBuilder(session, rootAgent)
.invocationId(resolvedInvocationId)
.runConfig(runConfig)
.userContent(
userContent != null
? userContent
: originalUserMessage(session, resolvedInvocationId))
.userContent(userContentFor(session, resolvedInvocationId, userContent))
.build();
// Rehydrate before routing, as Python does, so a root that checkpointed re-enters itself.
rootContext.populateInvocationAgentStates();
Expand Down Expand Up @@ -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();
}

/**
Expand Down Expand Up @@ -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)
Expand Down
Loading
Loading