Skip to content

feat: support pausing and resuming a session invocation #1473

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

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: support pausing and resuming a session invocation #1473

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
49 changes: 45 additions & 4 deletions core/src/main/java/com/google/adk/agents/BaseAgent.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,13 +16,15 @@

package com.google.adk.agents;

import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Strings.isNullOrEmpty;
import static com.google.common.collect.ImmutableList.toImmutableList;
import static java.lang.String.format;

import com.google.adk.agents.Callbacks.AfterAgentCallback;
import com.google.adk.agents.Callbacks.BeforeAgentCallback;
import com.google.adk.events.Event;
import com.google.adk.events.EventActions;
import com.google.adk.plugins.Plugin;
import com.google.adk.telemetry.Instrumentation;
import com.google.adk.telemetry.Instrumentation.AgentInvocation;
Expand All @@ -38,6 +40,7 @@
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.function.Function;
import java.util.regex.Pattern;
Expand Down Expand Up @@ -135,10 +138,8 @@ private static void validateAgentName(String name) {
throw new IllegalArgumentException(
format("Agent name '%s' does not match regex '%s'.", name, IDENTIFIER_REGEX));
}
if (name.equals(Role.USER)) {
throw new IllegalArgumentException(
"Agent name cannot be 'user'; reserved for end-user input.");
}
checkArgument(
!name.equals(Role.USER), "Agent name cannot be 'user'; reserved for end-user input.");
}

/**
Expand Down Expand Up @@ -459,6 +460,46 @@ public Flowable<Event> runLive(InvocationContext parentContext) {
return run(parentContext, this::runLiveImpl);
}

/**
* Records this agent's end-of-agent checkpoint and returns it as a single-event stream. Recording
* it via {@link InvocationContext#setAgentState} clears this agent's state and marks it finished,
* so a later run can skip it.
*
* @param context Current invocation context.
* @return a stream of the single {@code endOfAgent = true} checkpoint event.
*/
final Flowable<Event> endOfAgentAndRecord(InvocationContext context) {
context.setAgentState(name(), /* agentState= */ null, /* endOfAgent= */ true);
return Flowable.just(checkpointEvent(context, EventActions.builder().endOfAgent(true).build()));
}

/**
* Records {@code agentState} for this agent and returns the matching checkpoint event as a
* single-event stream. Recording it via {@link InvocationContext#setAgentState} (not yet ended)
* lets a later run resume at the right point.
*
* @param context Current invocation context.
* @param agentState The serialized agent state to persist.
* @return a stream of the single checkpoint event carrying {@code agentState}.
*/
final Flowable<Event> checkpointAndRecord(
InvocationContext context, Map<String, Object> agentState) {
context.setAgentState(name(), agentState, /* endOfAgent= */ false);
return Flowable.just(
checkpointEvent(context, EventActions.builder().agentState(agentState).build()));
}

/** Builds a resumability checkpoint event authored by this agent carrying {@code actions}. */
private Event checkpointEvent(InvocationContext context, EventActions actions) {
return Event.builder()
.id(Event.generateEventId())
.invocationId(context.invocationId())
.author(name())
.branch(context.branch().orElse(null))
.actions(actions)
.build();
}

/**
* Agent-specific asynchronous logic.
*
Expand Down
207 changes: 202 additions & 5 deletions core/src/main/java/com/google/adk/agents/InvocationContext.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,25 +20,36 @@

import com.google.adk.apps.ResumabilityConfig;
import com.google.adk.artifacts.BaseArtifactService;
import com.google.adk.events.Event;
import com.google.adk.memory.BaseMemoryService;
import com.google.adk.models.LlmCallsLimitExceededException;
import com.google.adk.plugins.Plugin;
import com.google.adk.plugins.PluginManager;
import com.google.adk.sessions.BaseSessionService;
import com.google.adk.sessions.Session;
import com.google.adk.summarizer.EventsCompactionConfig;
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 com.google.genai.types.Content;
import com.google.genai.types.FunctionCall;
import com.google.genai.types.FunctionResponse;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;
import org.jspecify.annotations.Nullable;

/** The context for an agent invocation. */
@SuppressWarnings("deprecation") // Plumbs the deprecated ResumabilityConfig.
public class InvocationContext {

private final BaseSessionService sessionService;
Expand All @@ -56,6 +67,10 @@ public class InvocationContext {
private final @Nullable ResumabilityConfig resumabilityConfig;
private final InvocationCostManager invocationCostManager;
private final Map<String, Object> callbackContextData;
// Resumability checkpoints, shared by reference across derived contexts so a sub-agent's
// checkpoint is visible to its parent and the runner.
private final Map<String, Map<String, Object>> agentStates;
private final Map<String, Boolean> endOfAgents;

@Nullable private String branch;
private BaseAgent agent;
Expand Down Expand Up @@ -83,6 +98,8 @@ protected InvocationContext(Builder builder) {
// invocation invocation so that Plugins can access the same data it during the invocation
// across all types of callbacks.
this.callbackContextData = builder.callbackContextData;
this.agentStates = builder.agentStates;
this.endOfAgents = builder.endOfAgents;
}

/** Returns a new {@link Builder} for creating {@link InvocationContext} instances. */
Expand Down Expand Up @@ -222,14 +239,188 @@ public Optional<ContextCacheConfig> contextCacheConfig() {
return Optional.ofNullable(contextCacheConfig);
}

/**
* Returns whether the current invocation is resumable. Mirrors Python ADK v1's {@code
* InvocationContext.is_resumable}.
*/
/** Returns whether the current invocation is resumable. */
public boolean isResumable() {
return resumabilityConfig != null && resumabilityConfig.isResumable();
}

/**
* Returns an unmodifiable view of the per-agent resumability checkpoint states for this
* invocation, keyed by agent name. The backing map is shared by reference across derived contexts
* within the invocation; mutate it only through {@link #setAgentState}.
*/
public Map<String, Map<String, Object>> agentStates() {
return Collections.unmodifiableMap(agentStates);
}

/**
* Returns an unmodifiable view of the per-agent end-of-agent flags for this invocation, keyed by
* agent name.
*/
public Map<String, Boolean> endOfAgents() {
return Collections.unmodifiableMap(endOfAgents);
}

/**
* Sets the checkpoint state of an agent explicitly. Does not implicitly initialize.
*
* @param agentName the agent whose state to set.
* @param agentState the serialized agent state to store; ignored when {@code endOfAgent} is true.
* @param endOfAgent when true, marks the agent finished and drops any stored state.
*/
void setAgentState(
String agentName, @Nullable Map<String, Object> agentState, boolean endOfAgent) {
if (endOfAgent) {
endOfAgents.put(agentName, true);
agentStates.remove(agentName);
} else if (agentState != null) {
// Store a read-only copy so agentStates() stays read-only (updates go through setAgentState).
// LinkedHashMap tolerates a null value that a deserialized (older-session) agentState may
// carry -- ImmutableMap.copyOf would reject it.
agentStates.put(agentName, Collections.unmodifiableMap(new LinkedHashMap<>(agentState)));
endOfAgents.put(agentName, false);
} else {
endOfAgents.remove(agentName);
agentStates.remove(agentName);
}
}

/** Recursively resets the checkpoint state of all sub-agents of the given agent. */
void resetSubAgentStates(String agentName) {
Optional<BaseAgent> target = agent.findAgent(agentName);
if (target.isEmpty()) {
return;
}
for (BaseAgent subAgent : target.get().subAgents()) {
setAgentState(subAgent.name(), /* agentState= */ null, /* endOfAgent= */ false);
resetSubAgentStates(subAgent.name());
}
}

/**
* Rehydrates {@link #agentStates()} and {@link #endOfAgents()} from the current invocation's
* history when this invocation is resumable. For each event carrying agent-state information,
* sets the authoring agent's checkpoint; for a non-workflow author that already produced content,
* seeds an empty state so it is treated as mid-run.
*/
public void populateInvocationAgentStates() {
if (!isResumable()) {
return;
}
for (Event event : events(/* currentInvocation= */ true, /* currentBranch= */ false)) {
String author = event.author();
if (author == null) {
continue;
}
Optional<Map<String, Object>> agentState = event.actions().agentState();
if (event.actions().endOfAgent()) {
endOfAgents.put(author, true);
agentStates.remove(author);
} else if (agentState.isPresent()) {
// setAgentState stores a null-tolerant read-only copy, so a deserialized (older-session)
// agentState carrying a null value does not crash resume.
setAgentState(author, agentState.get(), /* endOfAgent= */ false);
} else if (!author.equals(Role.USER)
&& event.content().isPresent()
&& !agentStates.containsKey(author)) {
agentStates.put(author, ImmutableMap.of());
endOfAgents.put(author, false);
}
}
}

/**
* Returns the current session's events, optionally filtered to the current invocation and/or the
* current branch. Reads the in-memory {@link Session#events()} list, which {@link
* BaseSessionService#appendEvent} keeps in sync. A {@code null}-branch event is visible on any
* branch.
*
* @param currentInvocation whether to filter to events from this invocation.
* @param currentBranch whether to filter to events on this branch (or with no branch).
*/
ImmutableList<Event> events(boolean currentInvocation, boolean currentBranch) {
List<Event> results = new ArrayList<>(session.events());
if (currentInvocation) {
results.removeIf(event -> !invocationId.equals(event.invocationId()));
}
if (currentBranch) {
results.removeIf(event -> event.branch().filter(b -> !b.equals(this.branch)).isPresent());
}
return ImmutableList.copyOf(results);
}

/**
* Returns whether to pause the invocation right after this event. Pausing (unlike ending) leaves
* the invocation resumable. Both conditions must hold: the app is {@link #isResumable()} and the
* event carries a long-running function call (including a synthetic {@code
* adk_request_confirmation} HITL request).
*/
boolean shouldPauseInvocation(Event event) {
if (!isResumable()) {
return false;
}
Set<String> longRunningIds = event.longRunningToolIds().orElse(ImmutableSet.of());
return event.functionCalls().stream()
.anyMatch(call -> call.id().filter(longRunningIds::contains).isPresent());
}

/**
* Returns whether a long-running call this invocation paused on is still unanswered. A resumed
* flow must not re-invoke the model while any long-running call it paused on lacks a response,
* including when resuming the paused call itself without an answer, so a partially answered set
* of parallel long-running calls keeps waiting while a fully answered one continues to a summary.
*/
boolean hasUnansweredPausedCall() {
List<Event> events = events(/* currentInvocation= */ true, /* currentBranch= */ true);
if (events.isEmpty()) {
return false;
}
Set<String> awaited = new HashSet<>();
for (Event event : events) {
if (!shouldPauseInvocation(event)) {
continue;
}
for (FunctionCall call : event.functionCalls()) {
call.id().ifPresent(awaited::add);
}
awaited.addAll(event.longRunningToolIds().orElse(ImmutableSet.of()));
}
if (awaited.isEmpty()) {
return false;
}
Set<String> answered = new HashSet<>();
for (Event event : events) {
for (FunctionResponse response : event.functionResponses()) {
response.id().ifPresent(answered::add);
}
}
return !answered.containsAll(awaited);
}

/**
* Finds the current-invocation event whose function call matches any function response id in
* {@code functionResponseEvent}, searching newest-first. Matching any id (not just the first)
* keeps parallel function responses resolvable when their calls interleave.
*/
Optional<Event> findMatchingFunctionCall(Event functionResponseEvent) {
Set<String> targetIds = new HashSet<>();
for (FunctionResponse response : functionResponseEvent.functionResponses()) {
response.id().ifPresent(targetIds::add);
}
if (targetIds.isEmpty()) {
return Optional.empty();
}
List<Event> events = events(/* currentInvocation= */ true, /* currentBranch= */ false);
for (int i = events.size() - 1; i >= 0; i--) {
for (FunctionCall call : events.get(i).functionCalls()) {
if (call.id().filter(targetIds::contains).isPresent()) {
return Optional.of(events.get(i));
}
}
}
return Optional.empty();
}

private static class InvocationCostManager {
private final AtomicInteger numberOfLlmCalls = new AtomicInteger(0);

Expand Down Expand Up @@ -289,6 +480,10 @@ private Builder(InvocationContext context) {
// invocation invocation so that Plugins can access the same data it during the invocation
// across all types of callbacks.
this.callbackContextData = context.callbackContextData;
// Shared by reference so a sub-agent's checkpoint is visible to its parent and the runner
// within one invocation.
this.agentStates = context.agentStates;
this.endOfAgents = context.endOfAgents;
}

private BaseSessionService sessionService;
Expand All @@ -309,6 +504,8 @@ private Builder(InvocationContext context) {
private @Nullable ResumabilityConfig resumabilityConfig;
private InvocationCostManager invocationCostManager = new InvocationCostManager();
private Map<String, Object> callbackContextData = new ConcurrentHashMap<>();
private Map<String, Map<String, Object>> agentStates = new ConcurrentHashMap<>();
private Map<String, Boolean> endOfAgents = new ConcurrentHashMap<>();

/**
* Sets the session service for managing session state.
Expand Down
Loading
Loading