diff --git a/core/src/main/java/com/google/adk/agents/BaseAgent.java b/core/src/main/java/com/google/adk/agents/BaseAgent.java index 0a3c550a4..36295c2b3 100644 --- a/core/src/main/java/com/google/adk/agents/BaseAgent.java +++ b/core/src/main/java/com/google/adk/agents/BaseAgent.java @@ -16,6 +16,7 @@ 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; @@ -23,6 +24,7 @@ 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; @@ -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; @@ -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."); } /** @@ -459,6 +460,46 @@ public Flowable 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 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 checkpointAndRecord( + InvocationContext context, Map 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. * diff --git a/core/src/main/java/com/google/adk/agents/InvocationContext.java b/core/src/main/java/com/google/adk/agents/InvocationContext.java index 456758b95..d95fdea21 100644 --- a/core/src/main/java/com/google/adk/agents/InvocationContext.java +++ b/core/src/main/java/com/google/adk/agents/InvocationContext.java @@ -20,6 +20,7 @@ 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; @@ -27,18 +28,28 @@ 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; @@ -56,6 +67,10 @@ public class InvocationContext { private final @Nullable ResumabilityConfig resumabilityConfig; private final InvocationCostManager invocationCostManager; private final Map 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> agentStates; + private final Map endOfAgents; @Nullable private String branch; private BaseAgent agent; @@ -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. */ @@ -222,14 +239,188 @@ public Optional 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> 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 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 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 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> 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 events(boolean currentInvocation, boolean currentBranch) { + List 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 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 events = events(/* currentInvocation= */ true, /* currentBranch= */ true); + if (events.isEmpty()) { + return false; + } + Set 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 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 findMatchingFunctionCall(Event functionResponseEvent) { + Set targetIds = new HashSet<>(); + for (FunctionResponse response : functionResponseEvent.functionResponses()) { + response.id().ifPresent(targetIds::add); + } + if (targetIds.isEmpty()) { + return Optional.empty(); + } + List 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); @@ -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; @@ -309,6 +504,8 @@ private Builder(InvocationContext context) { private @Nullable ResumabilityConfig resumabilityConfig; private InvocationCostManager invocationCostManager = new InvocationCostManager(); private Map callbackContextData = new ConcurrentHashMap<>(); + private Map> agentStates = new ConcurrentHashMap<>(); + private Map endOfAgents = new ConcurrentHashMap<>(); /** * Sets the session service for managing session state. 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 fa754e0c0..080272afb 100644 --- a/core/src/main/java/com/google/adk/agents/LlmAgent.java +++ b/core/src/main/java/com/google/adk/agents/LlmAgent.java @@ -55,6 +55,7 @@ import com.google.adk.tools.BaseToolset; import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableList; +import com.google.common.collect.Iterables; import com.google.errorprone.annotations.CanIgnoreReturnValue; import com.google.genai.types.Content; import com.google.genai.types.GenerateContentConfig; @@ -70,6 +71,7 @@ import java.util.Objects; import java.util.Optional; import java.util.concurrent.Executor; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.function.Function; import org.jspecify.annotations.Nullable; import org.slf4j.Logger; @@ -660,7 +662,133 @@ private static boolean isThought(Part part) { @Override protected Flowable runAsyncImpl(InvocationContext invocationContext) { - return llmFlow.run(invocationContext).doOnNext(this::maybeSaveOutputToState); + if (!invocationContext.isResumable()) { + return llmFlow.run(invocationContext).doOnNext(this::maybeSaveOutputToState); + } + return Flowable.defer( + () -> { + // Resumed after a transfer: continue the transferred sub-agent instead of re-invoking + // the model, then mark this agent done -- unless the sub-agent pauses again, so a later + // turn can resume that pause. + if (invocationContext.agentStates().containsKey(name())) { + Optional resumeTarget = findSubAgentToResume(invocationContext); + if (resumeTarget.isPresent()) { + AtomicBoolean resumePaused = new AtomicBoolean(false); + return resumeTarget + .get() + .runAsync(invocationContext) + .doOnNext( + event -> { + if (invocationContext.shouldPauseInvocation(event)) { + resumePaused.set(true); + } + }) + .concatWith( + Flowable.defer( + () -> { + if (resumePaused.get()) { + return Flowable.empty(); + } + return endOfAgentAndRecord(invocationContext); + })); + } + } + // Don't re-invoke the model while any paused long-running call is unanswered. + if (invocationContext.hasUnansweredPausedCall()) { + return Flowable.empty(); + } + // Normal path: emit an end-of-agent checkpoint on completion so a later run can skip + // this agent, unless it paused on a long-running call (then suppress it so it can + // resume). + Flowable events = + llmFlow.run(invocationContext).doOnNext(this::maybeSaveOutputToState); + AtomicBoolean paused = new AtomicBoolean(false); + AtomicBoolean transferred = new AtomicBoolean(false); + return events + .doOnNext( + event -> { + if (invocationContext.shouldPauseInvocation(event)) { + paused.set(true); + } + }) + .concatMap( + event -> { + // On a transfer this agent authored, close this agent here -- before the + // transferred-to sub-agent runs -- so its checkpoint marks it done and a later + // turn resumes at the sub-agent, not the finished root. + if (transferTargetFrom(event).isPresent()) { + transferred.set(true); + return Flowable.just(event) + .concatWith(endOfAgentAndRecord(invocationContext)); + } + return Flowable.just(event); + }) + .concatWith( + Flowable.defer( + () -> { + if (paused.get() || transferred.get()) { + return Flowable.empty(); + } + return endOfAgentAndRecord(invocationContext); + })); + }); + } + + /** + * Returns the agent this agent transferred to in {@code event} (an event this agent authored that + * carries a transfer to a different agent), or empty when {@code event} is not such a transfer. + */ + private Optional transferTargetFrom(Event event) { + if (name().equals(event.author())) { + return event + .actions() + .transferToAgent() + .filter(target -> !target.equals(name())) + .flatMap(target -> rootAgent().findAgent(target)); + } + return Optional.empty(); + } + + /** + * When this agent is being resumed, returns the sub-agent it had transferred to (so the resume + * continues that sub-agent), or empty when this agent should continue itself. + */ + private Optional findSubAgentToResume(InvocationContext context) { + List events = context.events(/* currentInvocation= */ true, /* currentBranch= */ true); + if (events.isEmpty()) { + return Optional.empty(); + } + Event lastEvent = Iterables.getLast(events); + if (name().equals(lastEvent.author())) { + return transferTargetFrom(lastEvent); + } + if (Objects.equals(lastEvent.author(), Role.USER)) { + // A plain-text resume message (no function response) is not a transfer resume: continue this + // agent rather than requiring a matching function call. + if (lastEvent.functionResponses().isEmpty()) { + return Optional.empty(); + } + // IAE (not ISE): an unresolvable resume surfaces through Runner.runAsync's IAE contract. + Event functionCallEvent = + context + .findMatchingFunctionCall(lastEvent) + .orElseThrow( + () -> + new IllegalArgumentException( + "No matching function call to resume agent " + + name() + + " from a function response.")); + if (name().equals(functionCallEvent.author())) { + return Optional.empty(); + } + } + for (int i = events.size() - 2; i >= 0; i--) { + Optional agent = transferTargetFrom(events.get(i)); + if (agent.isPresent()) { + return agent; + } + } + return Optional.empty(); } @Override diff --git a/core/src/main/java/com/google/adk/agents/LoopAgent.java b/core/src/main/java/com/google/adk/agents/LoopAgent.java index 19fd4c497..20011767c 100644 --- a/core/src/main/java/com/google/adk/agents/LoopAgent.java +++ b/core/src/main/java/com/google/adk/agents/LoopAgent.java @@ -18,9 +18,11 @@ import com.google.adk.agents.ConfigAgentUtils.ConfigurationException; import com.google.adk.events.Event; +import com.google.common.collect.ImmutableMap; import com.google.errorprone.annotations.CanIgnoreReturnValue; import io.reactivex.rxjava3.core.Flowable; import java.util.List; +import java.util.Map; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import org.jspecify.annotations.Nullable; @@ -149,29 +151,109 @@ protected Flowable runAsyncImpl(InvocationContext invocationContext) { .takeUntil(LoopAgent::hasEscalateAction); } - // Resumable: stop looping once a sub-agent emits a pending long-running call (e.g. HITL), - // matching Python ADK v1 and avoiding a runaway loop. The current sub-agent still finishes; - // resuming into the paused iteration needs persisted state (future work). - AtomicBoolean paused = new AtomicBoolean(false); - AtomicInteger timesLooped = new AtomicInteger(0); - return Flowable.fromIterable(subAgents) - .concatMap( - subAgent -> - paused.get() - ? Flowable.empty() - : subAgent - .runAsync(invocationContext) - .doOnNext( - event -> { - if (WorkflowAgentResumption.hasPendingLongRunningCall(event)) { - paused.set(true); - } - })) - .repeatUntil( - () -> - paused.get() - || (maxIterations != null && timesLooped.incrementAndGet() >= maxIterations)) - .takeUntil(LoopAgent::hasEscalateAction); + // Resumable: checkpoint {current_sub_agent, times_looped} before each sub-agent, resume into + // the checkpointed iteration, pause (not end) on a long-running call, and reset sub-agent + // state between iterations. + return Flowable.defer( + () -> { + Map state = invocationContext.agentStates().get(name()); + String startSubAgentName = + state != null && state.get(WorkflowAgentStates.CURRENT_SUB_AGENT) instanceof String s + ? s + : null; + int startTimesLooped = + state != null && state.get(WorkflowAgentStates.TIMES_LOOPED) instanceof Number n + ? n.intValue() + : 0; + int startIndex = + WorkflowAgentStates.findIndexForResumption(subAgents, startSubAgentName, logger); + LoopState loopState = new LoopState(startSubAgentName != null, startTimesLooped); + return runLoopIteration(invocationContext, subAgents, startIndex, loopState); + }); + } + + /** Mutable state shared across the iterations of one resumable {@link LoopAgent} run. */ + private static final class LoopState { + /** True until the sub-agent being resumed into has run; that sub-agent skips its checkpoint. */ + final AtomicBoolean resuming; + + final AtomicInteger timesLooped; + final AtomicBoolean shouldExit = new AtomicBoolean(false); + final AtomicBoolean paused = new AtomicBoolean(false); + + LoopState(boolean resuming, int timesLooped) { + this.resuming = new AtomicBoolean(resuming); + this.timesLooped = new AtomicInteger(timesLooped); + } + } + + /** + * Runs one loop iteration over the sub-agents from {@code startIndex}, then either recurses for + * the next iteration or terminates (emitting end-of-agent unless paused). {@code state} carries + * the loop's mutable state across iterations. + */ + private Flowable runLoopIteration( + InvocationContext context, + List subAgents, + int startIndex, + LoopState state) { + return Flowable.defer( + () -> { + // Iteration cap, checked in one place before each iteration: this covers both a resume + // that starts already at/over the cap and the transition after an iteration completes. + if (maxIterations != null && state.timesLooped.get() >= maxIterations) { + return endOfAgentAndRecord(context); + } + Flowable iteration = + Flowable.fromIterable(subAgents.subList(startIndex, subAgents.size())) + .concatMap( + subAgent -> + Flowable.defer( + () -> { + if (state.shouldExit.get() || state.paused.get()) { + return Flowable.empty(); + } + Flowable checkpoint = Flowable.empty(); + if (!state.resuming.getAndSet(false)) { + ImmutableMap subState = + ImmutableMap.of( + WorkflowAgentStates.CURRENT_SUB_AGENT, subAgent.name(), + WorkflowAgentStates.TIMES_LOOPED, + state.timesLooped.get()); + checkpoint = checkpointAndRecord(context, subState); + } + Flowable run = + subAgent + .runAsync(context) + .doOnNext( + event -> { + if (hasEscalateAction(event)) { + state.shouldExit.set(true); + } + if (context.shouldPauseInvocation(event)) { + state.paused.set(true); + } + }); + return checkpoint.concatWith(run); + })); + return iteration.concatWith( + Flowable.defer( + () -> { + // Pause takes precedence over escalation-exit so a long-running pause stays + // resumable. + if (state.paused.get()) { + return Flowable.empty(); + } + if (state.shouldExit.get()) { + return endOfAgentAndRecord(context); + } + state.timesLooped.incrementAndGet(); + context.resetSubAgentStates(name()); + // A fresh iteration restarts at the first sub-agent (state.resuming is already + // false here). The cap is re-checked at the top of the next iteration. + return runLoopIteration(context, subAgents, /* startIndex= */ 0, state); + })); + }); } @Override diff --git a/core/src/main/java/com/google/adk/agents/ParallelAgent.java b/core/src/main/java/com/google/adk/agents/ParallelAgent.java index e1382a317..bf56bd219 100644 --- a/core/src/main/java/com/google/adk/agents/ParallelAgent.java +++ b/core/src/main/java/com/google/adk/agents/ParallelAgent.java @@ -19,12 +19,14 @@ import com.google.adk.agents.ConfigAgentUtils.ConfigurationException; import com.google.adk.events.Event; +import com.google.common.collect.ImmutableMap; import com.google.errorprone.annotations.CanIgnoreReturnValue; import io.reactivex.rxjava3.core.Flowable; import io.reactivex.rxjava3.core.Scheduler; import io.reactivex.rxjava3.schedulers.Schedulers; import java.util.ArrayList; import java.util.List; +import java.util.concurrent.atomic.AtomicBoolean; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -175,13 +177,79 @@ protected Flowable runAsyncImpl(InvocationContext invocationContext) { return Flowable.empty(); } - var updatedInvocationContext = setBranchForCurrentAgent(this, invocationContext); - List> agentFlowables = new ArrayList<>(); - for (BaseAgent subAgent : currentSubAgents) { - agentFlowables.add(subAgent.runAsync(updatedInvocationContext).subscribeOn(scheduler)); + if (!invocationContext.isResumable()) { + var updatedInvocationContext = setBranchForCurrentAgent(this, invocationContext); + List> agentFlowables = new ArrayList<>(); + for (BaseAgent subAgent : currentSubAgents) { + agentFlowables.add(subAgent.runAsync(updatedInvocationContext).subscribeOn(scheduler)); + } + return Flowable.merge(agentFlowables) + .takeUntil((Event event) -> event.actions().escalate().orElse(false)); } - return Flowable.merge(agentFlowables) - .takeUntil((Event event) -> event.actions().escalate().orElse(false)); + + // Resumable: skip completed branches, checkpoint that this agent started, pause (without + // ending) if any branch pauses, and end only once every active branch finished. + return Flowable.defer( + () -> { + List activeSubAgents = new ArrayList<>(); + for (BaseAgent subAgent : currentSubAgents) { + if (!invocationContext.endOfAgents().getOrDefault(subAgent.name(), false)) { + activeSubAgents.add(subAgent); + } + } + + Flowable initialCheckpoint = Flowable.empty(); + if (!invocationContext.agentStates().containsKey(name())) { + initialCheckpoint = checkpointAndRecord(invocationContext, ImmutableMap.of()); + } + + var updatedInvocationContext = setBranchForCurrentAgent(this, invocationContext); + AtomicBoolean paused = new AtomicBoolean(false); + AtomicBoolean escalated = new AtomicBoolean(false); + List> agentFlowables = new ArrayList<>(); + for (BaseAgent subAgent : activeSubAgents) { + agentFlowables.add( + subAgent + .runAsync(updatedInvocationContext) + .subscribeOn(scheduler) + .doOnNext( + event -> { + if (invocationContext.shouldPauseInvocation(event)) { + paused.set(true); + } + if (event.actions().escalate().orElse(false)) { + escalated.set(true); + } + })); + } + Flowable merged = + Flowable.merge(agentFlowables) + .takeUntil((Event event) -> event.actions().escalate().orElse(false)); + + return initialCheckpoint + .concatWith(merged) + .concatWith( + Flowable.defer( + () -> { + if (paused.get()) { + return Flowable.empty(); + } + // A sub-agent escalation ends this agent even if other branches did not + // finish; otherwise it ends once every active branch finished (a custom + // BaseAgent that never records endOfAgent may not reach the latter). + boolean allEnded = + activeSubAgents.stream() + .allMatch( + a -> + invocationContext + .endOfAgents() + .getOrDefault(a.name(), false)); + if (escalated.get() || allEnded) { + return endOfAgentAndRecord(invocationContext); + } + return Flowable.empty(); + })); + }); } /** diff --git a/core/src/main/java/com/google/adk/agents/SequentialAgent.java b/core/src/main/java/com/google/adk/agents/SequentialAgent.java index 963c3d109..fb30056fd 100644 --- a/core/src/main/java/com/google/adk/agents/SequentialAgent.java +++ b/core/src/main/java/com/google/adk/agents/SequentialAgent.java @@ -17,8 +17,11 @@ import com.google.adk.agents.ConfigAgentUtils.ConfigurationException; import com.google.adk.events.Event; +import com.google.common.collect.ImmutableMap; import io.reactivex.rxjava3.core.Flowable; import java.util.List; +import java.util.Map; +import java.util.Optional; import java.util.concurrent.atomic.AtomicBoolean; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -92,8 +95,7 @@ public static Builder builder() { * *

When resumability is enabled, on resume execution fast-forwards to the sub-agent being * resumed (completed ones are not re-run) and pauses on a pending long-running call; when - * disabled, sub-agents simply run in order (matches Python ADK v1 with resumability off). - * Temporary, event-based. + * disabled, sub-agents run in order. * * @param invocationContext Invocation context. * @return Flowable emitting events from sub-agents. @@ -108,22 +110,61 @@ protected Flowable runAsyncImpl(InvocationContext invocationContext) { return Flowable.fromIterable(subAgents) .concatMap(subAgent -> subAgent.runAsync(invocationContext)); } - int startIndex = - WorkflowAgentResumption.resumeSubAgentIndex(invocationContext, subAgents).orElse(0); + // Resumable: checkpoint each sub-agent before it runs, fast-forward to the checkpoint on + // resume, and pause (without ending) on a long-running call. + Map state = invocationContext.agentStates().get(name()); + String startSubAgentName = + state != null && state.get(WorkflowAgentStates.CURRENT_SUB_AGENT) instanceof String s + ? s + : null; + int startIndex; + boolean isResuming; + if (startSubAgentName != null) { + startIndex = WorkflowAgentStates.findIndexForResumption(subAgents, startSubAgentName, logger); + isResuming = true; + } else { + // Back-compat: a session paused before checkpoints existed has no agentState; reconstruct + // the resume point from history so it still fast-forwards past completed sub-agents. + Optional reconstructed = + WorkflowAgentResumption.resumeSubAgentIndex(invocationContext, subAgents); + startIndex = reconstructed.orElse(0); + isResuming = reconstructed.isPresent(); + } AtomicBoolean paused = new AtomicBoolean(false); + AtomicBoolean resuming = new AtomicBoolean(isResuming); return Flowable.fromIterable(subAgents.subList(startIndex, subAgents.size())) .concatMap( subAgent -> - paused.get() - ? Flowable.empty() - : subAgent - .runAsync(invocationContext) - .doOnNext( - event -> { - if (WorkflowAgentResumption.hasPendingLongRunningCall(event)) { - paused.set(true); - } - })); + Flowable.defer( + () -> { + if (paused.get()) { + return Flowable.empty(); + } + Flowable checkpoint = Flowable.empty(); + if (!resuming.getAndSet(false)) { + ImmutableMap subState = + ImmutableMap.of(WorkflowAgentStates.CURRENT_SUB_AGENT, subAgent.name()); + checkpoint = checkpointAndRecord(invocationContext, subState); + } + Flowable run = + subAgent + .runAsync(invocationContext) + .doOnNext( + event -> { + if (invocationContext.shouldPauseInvocation(event)) { + paused.set(true); + } + }); + return checkpoint.concatWith(run); + })) + .concatWith( + Flowable.defer( + () -> { + if (paused.get()) { + return Flowable.empty(); + } + return endOfAgentAndRecord(invocationContext); + })); } /** diff --git a/core/src/main/java/com/google/adk/agents/WorkflowAgentResumption.java b/core/src/main/java/com/google/adk/agents/WorkflowAgentResumption.java index 2bff47803..a0a901b08 100644 --- a/core/src/main/java/com/google/adk/agents/WorkflowAgentResumption.java +++ b/core/src/main/java/com/google/adk/agents/WorkflowAgentResumption.java @@ -22,8 +22,9 @@ import java.util.Optional; /** - * Helpers for resuming workflow agents from session events. Temporary until session resumption - * (persisted agent state) is available. + * Back-compat helper: reconstructs a workflow agent's resume point from session events for + * invocations paused before durable agent-state checkpoints existed (used from {@link + * SequentialAgent} when no checkpoint state is present). */ final class WorkflowAgentResumption { @@ -48,12 +49,5 @@ static Optional resumeSubAgentIndex( return Optional.empty(); } - /** - * Whether the event emits a long-running call still awaiting a response (e.g. a HITL request). - */ - static boolean hasPendingLongRunningCall(Event event) { - return Functions.hasPendingLongRunningCall(event); - } - private WorkflowAgentResumption() {} } diff --git a/core/src/main/java/com/google/adk/agents/WorkflowAgentStates.java b/core/src/main/java/com/google/adk/agents/WorkflowAgentStates.java new file mode 100644 index 000000000..bbf08cb35 --- /dev/null +++ b/core/src/main/java/com/google/adk/agents/WorkflowAgentStates.java @@ -0,0 +1,55 @@ +/* + * Copyright 2025 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.adk.agents; + +import java.util.List; +import org.jspecify.annotations.Nullable; +import org.slf4j.Logger; + +/** + * Wire-format keys and helpers for workflow-agent resumability checkpoints. The keys match Python + * and Kotlin ADK so persisted state is portable across languages. + */ +final class WorkflowAgentStates { + + /** Key holding the name of the current/next sub-agent in a Sequential or Loop checkpoint. */ + static final String CURRENT_SUB_AGENT = "current_sub_agent"; + + /** Key holding the completed-iteration count in a Loop checkpoint. */ + static final String TIMES_LOOPED = "times_looped"; + + /** + * Returns the index of the sub-agent to resume from by name, or 0 when the name is null or (with + * a warning) when it is no longer present in the sub-agents list. + */ + static int findIndexForResumption( + List subAgents, @Nullable String agentName, Logger logger) { + if (agentName == null) { + return 0; + } + for (int i = 0; i < subAgents.size(); i++) { + if (agentName.equals(subAgents.get(i).name())) { + return i; + } + } + // Agent names are developer-assigned identifiers, not user data, so log the missing name. + logger.warn("Restored sub-agent '{}' not found; resuming from index 0.", agentName); + return 0; + } + + private WorkflowAgentStates() {} +} diff --git a/core/src/main/java/com/google/adk/apps/App.java b/core/src/main/java/com/google/adk/apps/App.java index 9120954b5..7a1b85878 100644 --- a/core/src/main/java/com/google/adk/apps/App.java +++ b/core/src/main/java/com/google/adk/apps/App.java @@ -19,6 +19,7 @@ import com.google.adk.agents.BaseAgent; import com.google.adk.agents.ContextCacheConfig; import com.google.adk.agents.Role; +import com.google.adk.annotations.Experimental; import com.google.adk.plugins.Plugin; import com.google.adk.summarizer.EventsCompactionConfig; import com.google.common.collect.ImmutableList; @@ -35,7 +36,6 @@ * and communication across all agents in the hierarchy. The {@code plugins} are application-wide * components that provide shared capabilities and services to the entire system. */ -@SuppressWarnings("deprecation") // Plumbs the deprecated ResumabilityConfig. public class App { private static final Pattern IDENTIFIER_PATTERN = Pattern.compile("[a-zA-Z_][a-zA-Z0-9_]*"); @@ -83,6 +83,7 @@ public ContextCacheConfig contextCacheConfig() { return contextCacheConfig; } + @Experimental public @Nullable ResumabilityConfig resumabilityConfig() { return resumabilityConfig; } @@ -132,14 +133,9 @@ public Builder contextCacheConfig(ContextCacheConfig contextCacheConfig) { return this; } - /** - * Sets the app resumability config. - * - * @deprecated See {@link ResumabilityConfig}: partial feature, full resumability not yet - * available. - */ + /** Sets the app resumability config. Experimental; see {@link ResumabilityConfig}. */ @CanIgnoreReturnValue - @Deprecated + @Experimental public Builder resumabilityConfig(ResumabilityConfig resumabilityConfig) { this.resumabilityConfig = resumabilityConfig; return this; diff --git a/core/src/main/java/com/google/adk/apps/ResumabilityConfig.java b/core/src/main/java/com/google/adk/apps/ResumabilityConfig.java index d6d0445d7..f2981908b 100644 --- a/core/src/main/java/com/google/adk/apps/ResumabilityConfig.java +++ b/core/src/main/java/com/google/adk/apps/ResumabilityConfig.java @@ -16,19 +16,18 @@ package com.google.adk.apps; +import com.google.adk.annotations.Experimental; import com.google.auto.value.AutoValue; import com.google.errorprone.annotations.CanIgnoreReturnValue; /** - * App resumability config, mirroring Python ADK v1's {@code ResumabilityConfig}: pause on a - * long-running call and resume from the last event. Applies to all agents in the app. + * App resumability config: pause on a long-running call and resume from the last event. Applies to + * all agents in the app. * - * @deprecated Partial feature: only event-reconstruction-based pause/resume for {@code - * SequentialAgent} is implemented. Full session resumability (persisted agent state, durable - * resume, other workflow agents) is not yet available. Forward-compatible: the same config will - * drive full resumability once it lands. + *

Experimental and not yet stable: resume is best-effort and at-least-once, so a resuming tool + * must be idempotent and any temporary in-memory state is lost on resumption. */ -@Deprecated +@Experimental @AutoValue public abstract class ResumabilityConfig { @@ -38,8 +37,8 @@ public abstract class ResumabilityConfig { /** * Whether a plain-text {@code runAsync} continuation -- a user message that is not a function * response -- resumes the last unfinished invocation instead of starting a new one. Off by - * default, matching Python ADK, where a plain-text {@code runAsync} always starts a new - * invocation and a paused invocation is resumed explicitly. + * default: a plain-text {@code runAsync} starts a new invocation and a paused invocation is + * resumed explicitly. * * @deprecated Back-compat shim for callers that deliver a resume as a plain-text turn. Migrate to * {@code Runner.runAsync(userId, sessionId, invocationId, message, runConfig, stateDelta)} diff --git a/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java b/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java index 91cc225f2..1811a760d 100644 --- a/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java +++ b/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java @@ -536,9 +536,8 @@ private Flowable run( return Flowable.empty(); } else if (invocationContext.isResumable() && Functions.hasPendingLongRunningCall(eventList)) { - // When resumable, a pending long-running call (e.g. HITL) pauses the flow - // instead of calling the model again, matching Python ADK v1 and avoiding a - // runaway re-issue loop. The disabled path is unchanged. + // Resumable: pause on an unanswered long-running call, but continue once it has + // a response. logger.debug("Pausing flow execution on a pending long-running call."); return Flowable.empty(); } else { diff --git a/core/src/main/java/com/google/adk/flows/llmflows/Functions.java b/core/src/main/java/com/google/adk/flows/llmflows/Functions.java index 2b5c07435..d169ae9b5 100644 --- a/core/src/main/java/com/google/adk/flows/llmflows/Functions.java +++ b/core/src/main/java/com/google/adk/flows/llmflows/Functions.java @@ -472,11 +472,32 @@ public static boolean hasPendingLongRunningCall(Event event) { } /** - * Returns whether the last one or two events hold a pending long-running call, meaning a - * resumable flow should pause instead of calling the model again. Mirrors Python ADK v1's - * flow-level pause check on {@code events[-1]} and {@code events[-2]}. + * Returns whether the last one or two events hold a long-running call still awaiting a response, + * meaning a resumable flow should pause instead of calling the model again. A response that + * resolves the call -- the tool's own same-turn value, or a later user-injected resume -- lets + * the flow continue and the model summarize; only a no-response return (null or empty result, + * which emits no function response) pauses. */ static boolean hasPendingLongRunningCall(List events) { + if (events.isEmpty()) { + return false; + } + Event last = Iterables.getLast(events); + if (events.size() >= 2 && !last.functionResponses().isEmpty()) { + Event pending = events.get(events.size() - 2); + Set longRunningIds = pending.longRunningToolIds().orElse(ImmutableSet.of()); + Set pausedIds = new HashSet<>(); + for (FunctionCall call : pending.functionCalls()) { + call.id().filter(longRunningIds::contains).ifPresent(pausedIds::add); + } + Set resolvedIds = new HashSet<>(); + for (FunctionResponse response : last.functionResponses()) { + response.id().ifPresent(resolvedIds::add); + } + if (!pausedIds.isEmpty() && resolvedIds.containsAll(pausedIds)) { + return false; + } + } int from = Math.max(0, events.size() - 2); for (int i = events.size() - 1; i >= from; i--) { if (hasPendingLongRunningCall(events.get(i))) { 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 5e23eb78c..1def18ace 100644 --- a/core/src/main/java/com/google/adk/runner/Runner.java +++ b/core/src/main/java/com/google/adk/runner/Runner.java @@ -17,6 +17,7 @@ package com.google.adk.runner; import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Preconditions.checkState; import com.google.adk.agents.ActiveStreamingTool; import com.google.adk.agents.BaseAgent; @@ -24,9 +25,11 @@ import com.google.adk.agents.InvocationContext; import com.google.adk.agents.LiveRequestQueue; import com.google.adk.agents.LlmAgent; +import com.google.adk.agents.LoopAgent; import com.google.adk.agents.Role; import com.google.adk.agents.RunConfig; import com.google.adk.agents.SequentialAgent; +import com.google.adk.annotations.Experimental; import com.google.adk.apps.App; import com.google.adk.apps.ResumabilityConfig; import com.google.adk.artifacts.BaseArtifactService; @@ -52,10 +55,13 @@ import com.google.adk.utils.CollectionUtils; import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; import com.google.common.collect.MapMaker; import com.google.errorprone.annotations.CanIgnoreReturnValue; import com.google.genai.types.AudioTranscriptionConfig; import com.google.genai.types.Content; +import com.google.genai.types.FunctionCall; +import com.google.genai.types.FunctionResponse; import com.google.genai.types.Modality; import com.google.genai.types.Part; import io.opentelemetry.api.trace.Span; @@ -69,15 +75,17 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; +import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; import org.jspecify.annotations.Nullable; /** The main class for the GenAI Agents runner. */ -@SuppressWarnings("deprecation") // Plumbs the deprecated ResumabilityConfig. +@SuppressWarnings("deprecation") // Reads the deprecated plainTextContinuationAutoResume shim. public class Runner { private final BaseAgent agent; private final String appName; @@ -359,6 +367,22 @@ private Single appendNewMessageToSession( InvocationContext invocationContext, boolean saveInputBlobsAsArtifacts, @Nullable Map stateDelta) { + return appendNewMessageToSession( + session, + newMessage, + invocationContext, + saveInputBlobsAsArtifacts, + stateDelta, + /* branch= */ null); + } + + private Single appendNewMessageToSession( + Session session, + Content newMessage, + InvocationContext invocationContext, + boolean saveInputBlobsAsArtifacts, + @Nullable Map stateDelta, + @Nullable String branch) { checkArgument(newMessage.parts().isPresent(), "No parts in the new_message."); Content messageToAppend = newMessage; @@ -393,6 +417,7 @@ private Single appendNewMessageToSession( .id(Event.generateEventId()) .invocationId(invocationContext.invocationId()) .author(Role.USER) + .branch(branch) .content(messageToAppend); // Add state delta if provided @@ -506,6 +531,52 @@ public Flowable runAsync(String userId, String sessionId, Content newMess return runAsync(userId, sessionId, newMessage, RunConfig.builder().build()); } + /** + * 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. + * + * @param userId the user id of the session. + * @param sessionId the session id. + * @param invocationId the invocation to resume; may be {@code null} when it can be inferred from + * {@code newMessage}. + * @param newMessage an optional message (typically a function response) to append before running. + * @param runConfig the run configuration. + * @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 IllegalArgumentException if the invocation cannot be resolved. + */ + @Experimental + public Flowable runAsync( + String userId, + String sessionId, + @Nullable String invocationId, + @Nullable Content newMessage, + RunConfig runConfig, + @Nullable Map stateDelta) { + checkState( + isResumable(), + "Resuming an invocation requires an App configured with a resumable ResumabilityConfig."); + return Flowable.defer( + () -> + this.sessionService + .getSession(appName, userId, sessionId, Optional.empty()) + .switchIfEmpty( + Single.error( + () -> + new IllegalArgumentException( + String.format( + "Session not found: %s for user %s", sessionId, userId)))) + .flatMapPublisher( + session -> + runResumableFromSession( + session, invocationId, newMessage, runConfig, stateDelta))) + .compose(Tracing.trace("invocation")); + } + /** * Runs the agent asynchronously using a provided Session object. * @@ -523,6 +594,21 @@ protected Flowable runAsyncImpl( Preconditions.checkNotNull(session, "session cannot be null"); Preconditions.checkNotNull(newMessage, "newMessage cannot be null"); Preconditions.checkNotNull(runConfig, "runConfig cannot be null"); + // When resumable, a message that resolves to an existing invocation (e.g. a function response + // to a paused call) resumes it; any other message starts a new invocation. Disabled: unchanged. + if (isResumable()) { + 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). */ + private Flowable runNewInvocation( + Session session, + Content newMessage, + RunConfig runConfig, + @Nullable Map stateDelta) { return Flowable.defer( () -> { Context capturedContext = Context.current(); @@ -559,12 +645,7 @@ protected Flowable runAsyncImpl( event -> runAgentWithUpdatedSession(initialContext, session, event, rootAgent) .compose(Tracing.withContext(capturedContext))) - .doOnError( - throwable -> - this.pluginManager - .runOnRunErrorCallback(initialContext, throwable) - .onErrorComplete() - .subscribe()); + .doOnError(throwable -> runOnRunError(initialContext, throwable)); }) .doOnError( throwable -> { @@ -598,27 +679,36 @@ private Flowable runAgentWithUpdatedSession( .userContent(event.content().orElseGet(Content::fromParts)) .build(); - // Call beforeRunCallback with updated session - Maybe beforeRunEvent = - this.pluginManager - .beforeRunCallback(contextWithUpdatedSession) - .map( - content -> - Event.builder() - .id(Event.generateEventId()) - .invocationId(contextWithUpdatedSession.invocationId()) - .author("model") - .content(content) - .build()); + // If beforeRunCallback returns content, emit it and skip agent. + Maybe beforeRunEvent = beforeRunEventFor(contextWithUpdatedSession); + Context capturedContext = Context.current(); + return executeAgentPipeline( + contextWithUpdatedSession, + updatedSession, + beforeRunEvent, + // TODO: remove this hack after deprecating runAsync with Session. + () -> copySessionStates(updatedSession, initialContext.session())) + .compose(Tracing.withContext(capturedContext)); + } + /** + * Runs {@code context.agent()} and drives the shared event pipeline both invocation paths use: + * persist each non-partial event (releasing the {@link PersistBarrier} step), run {@code + * onEachPersisted} if given, fire the {@code onEvent} plugin callback, then run the after-run and + * compaction brackets; a {@code beforeRunEvent} short-circuits the agent run. + */ + private Flowable executeAgentPipeline( + InvocationContext context, + Session sessionToPersist, + Maybe beforeRunEvent, + @Nullable Runnable onEachPersisted) { // Let BaseLlmFlow block each step until this Runner has persisted the prior step's events. - PersistBarrier.enable(contextWithUpdatedSession); + PersistBarrier.enable(context); - // Agent execution Flowable agentEvents = - contextWithUpdatedSession + context .agent() - .runAsync(contextWithUpdatedSession) + .runAsync(context) .concatMap( agentEvent -> { // Mirror ADK Python (runners.py): partial events are streamed to the caller but @@ -628,39 +718,32 @@ private Flowable runAgentWithUpdatedSession( Single persistStep = agentEvent.partial().orElse(false) ? Single.just(agentEvent) - : this.sessionService.appendEvent(updatedSession, agentEvent); + : this.sessionService.appendEvent(sessionToPersist, agentEvent); return persistStep // Release (or fail) BaseLlmFlow's wait for this step; the Runner stays the // sole appendEvent caller (see PersistBarrier). .doOnSuccess( - unusedEvent -> - PersistBarrier.markPersisted( - contextWithUpdatedSession, agentEvent.id())) + unusedEvent -> PersistBarrier.markPersisted(context, agentEvent.id())) .doOnError( - error -> - PersistBarrier.markFailed( - contextWithUpdatedSession, agentEvent.id(), error)) + error -> PersistBarrier.markFailed(context, agentEvent.id(), error)) .flatMap( registeredEvent -> { - // TODO: remove this hack after deprecating runAsync with Session. - copySessionStates(updatedSession, initialContext.session()); - return contextWithUpdatedSession + if (onEachPersisted != null) { + onEachPersisted.run(); + } + return context .pluginManager() - .onEventCallback(contextWithUpdatedSession, registeredEvent) + .onEventCallback(context, registeredEvent) .defaultIfEmpty(registeredEvent); }) .toFlowable(); }); - // If beforeRunCallback returns content, emit it and skip agent - Context capturedContext = Context.current(); return beforeRunEvent .toFlowable() .switchIfEmpty(agentEvents) - .concatWith( - Completable.defer(() -> pluginManager.afterRunCallback(contextWithUpdatedSession))) - .concatWith(Completable.defer(() -> compactEvents(updatedSession))) - .compose(Tracing.withContext(capturedContext)); + .concatWith(Completable.defer(() -> pluginManager.afterRunCallback(context))) + .concatWith(Completable.defer(() -> compactEvents(sessionToPersist))); } private Completable compactEvents(Session session) { @@ -671,6 +754,334 @@ private Completable compactEvents(Session session) { .orElseGet(Completable::complete); } + /** + * The optional before-run event: when a before-run callback returns content, wrap it as a model + * event that short-circuits the agent run. Both invocation paths use this. + */ + private Maybe beforeRunEventFor(InvocationContext context) { + return this.pluginManager + .beforeRunCallback(context) + .map( + content -> + Event.builder() + .id(Event.generateEventId()) + .invocationId(context.invocationId()) + .author("model") + .content(content) + .build()); + } + + /** Fires the on-run-error plugin callback; the run paths share this error handler. */ + private void runOnRunError(InvocationContext context, Throwable throwable) { + this.pluginManager.runOnRunErrorCallback(context, throwable).onErrorComplete().subscribe(); + } + + /** + * Resumes an existing invocation when one resolves from {@code providedInvocationId} or a + * function response in {@code newMessage}; otherwise starts a new invocation. Requires + * resumability. + */ + private Flowable runResumableFromSession( + Session session, + @Nullable String providedInvocationId, + @Nullable Content newMessage, + RunConfig runConfig, + @Nullable Map stateDelta) { + return Flowable.defer( + () -> { + // A function response whose id matches no call in history cannot resume any invocation + // (it would feed the model an orphan response); reject it before any invocation-id + // fallback or auto-resume, matching Python and Kotlin. + if (newMessage != null + && hasFunctionResponse(newMessage) + && matchingFunctionCallEvent(session, newMessage).isEmpty()) { + return Flowable.error( + new IllegalArgumentException( + "No matching function call for the function response in the resume message.")); + } + String resolvedInvocationId = + resolveInvocationId(session, newMessage, providedInvocationId); + if (resolvedInvocationId == null + && newMessage != null + && plainTextContinuationAutoResumes()) { + // Deprecated opt-in: resume the last unfinished invocation on a plain-text + // continuation (legacy behavior). + resolvedInvocationId = lastUnfinishedInvocationId(session); + } + if (resolvedInvocationId == null) { + if (newMessage == null) { + return Flowable.error( + new IllegalArgumentException( + "No new message provided and no resumable invocation to resume.")); + } + return runNewInvocation(session, newMessage, runConfig, stateDelta); + } + if (!sessionHasEventsForInvocation(session, resolvedInvocationId)) { + // Resume was requested for an invocation the session has no events for. + return Flowable.error( + new IllegalArgumentException("No events to resume for the requested invocation.")); + } + return resumeCore(session, resolvedInvocationId, newMessage, runConfig, stateDelta); + }); + } + + /** + * Runs an existing invocation on the given session: optionally appends {@code newMessage}, + * rehydrates agent checkpoints, skips a completed invocation, and runs the resolved agent under + * the resumed invocation id. + */ + private Flowable resumeCore( + Session session, + String resolvedInvocationId, + @Nullable Content newMessage, + RunConfig runConfig, + @Nullable Map stateDelta) { + return Flowable.defer( + () -> { + Context capturedContext = Context.current(); + if (stateDelta != null && !stateDelta.isEmpty()) { + stateDelta.forEach((key, value) -> session.state().put(key, value)); + } + + // Context for the pre-run plugin callbacks; the agent to run is (re)resolved after any + // append in runResumedAgent. + InvocationContext initialContext = + newInvocationContextBuilder(session) + .invocationId(resolvedInvocationId) + .runConfig(runConfig) + .userContent(newMessage == null ? Content.fromParts() : newMessage) + .build(); + + Flowable events; + if (newMessage != null) { + // Run the same on-user-message plugin callback as the new-invocation path, then append + // the function-response message under the resumed invocation first, inheriting the + // branch of the call it answers, so routing and rehydration see it. + String branch = + matchingFunctionCallEvent(session, newMessage).flatMap(Event::branch).orElse(null); + events = + this.pluginManager + .onUserMessageCallback(initialContext, newMessage) + .compose(Tracing.withContext(capturedContext)) + .defaultIfEmpty(newMessage) + .flatMap( + content -> + appendNewMessageToSession( + session, + content, + initialContext, + runConfig.saveInputBlobsAsArtifacts(), + stateDelta, + branch)) + .flatMapPublisher( + userEvent -> + runResumedAgent( + session, + resolvedInvocationId, + userEvent.content().orElse(null), + runConfig)); + } else { + events = + runResumedAgent(session, resolvedInvocationId, /* userContent= */ null, runConfig); + } + + return events + .doOnError(throwable -> runOnRunError(initialContext, throwable)) + .compose(Tracing.withContext(capturedContext)); + }); + } + + /** + * Runs the resolved agent for a resumed invocation with the same before-run / after-run plugin + * bracket as the new-invocation path. Rehydrates checkpoints, skips a completed invocation, and + * persists each event. + */ + private Flowable runResumedAgent( + Session session, + String resolvedInvocationId, + @Nullable Content userContent, + RunConfig runConfig) { + return Flowable.defer( + () -> { + // Build the resumed context on the resolved agent's parent branch (see + // resumeParentBranch). + BaseAgent resumeAgent = findAgentToRun(session, this.agent); + InvocationContext context = + newInvocationContextBuilder(session) + .invocationId(resolvedInvocationId) + .branch(resumeParentBranch(session, resolvedInvocationId, resumeAgent)) + .runConfig(runConfig) + .userContent(userContent == null ? Content.fromParts() : userContent) + .build(); + context.populateInvocationAgentStates(); + + // No-op guard: a completed invocation (its active agent already finished) is not re-run. + if (context.endOfAgents().getOrDefault(context.agent().name(), false)) { + return Flowable.empty(); + } + + // before_run may short-circuit the run with a model event, as on the new-invocation path. + Maybe beforeRunEvent = beforeRunEventFor(context); + + return executeAgentPipeline( + context, session, beforeRunEvent, /* onEachPersisted= */ null); + }); + } + + /** + * Branch to seed a resumed context with so {@code resumeAgent} runs under the same branch it + * originally did. Returns the parent branch (the resolved agent's most recent event branch minus + * its own trailing name segment, which {@link BaseAgent#runAsync} re-appends), or {@code null} + * for the root branch. Non-null only for an agent nested under a {@link ParallelAgent}. + */ + private static @Nullable String resumeParentBranch( + Session session, String invocationId, BaseAgent resumeAgent) { + List events = session.events(); + for (int i = events.size() - 1; i >= 0; i--) { + Event event = events.get(i); + if (invocationId.equals(event.invocationId()) + && resumeAgent.name().equals(event.author()) + && event.branch().isPresent()) { + String branch = event.branch().get(); + String ownSegment = "." + resumeAgent.name(); + if (branch.endsWith(ownSegment)) { + String parent = branch.substring(0, branch.length() - ownSegment.length()); + return parent.isEmpty() ? null : parent; + } + return branch.equals(resumeAgent.name()) ? null : branch; + } + } + return null; + } + + /** + * Resolves which invocation a request targets: the invocation that issued the function call + * matching {@code newMessage}'s function response, else the caller-supplied {@code invocationId}. + * Returns {@code null} when neither applies (a fresh message starts a new invocation). + */ + private static @Nullable String resolveInvocationId( + Session session, @Nullable Content newMessage, @Nullable String invocationId) { + if (newMessage != null) { + return matchingFunctionCallEvent(session, newMessage) + .map(Event::invocationId) + .orElse(invocationId); + } + return invocationId; + } + + /** + * Returns the id of the session's most recent invocation that has not finished (its active agent + * has not emitted an end-of-agent checkpoint), or {@code null} when the last invocation already + * finished, when there is none, or when an uncheckpointed session has no invocation paused on an + * unanswered long-running call. Used only by the deprecated plain-text auto-resume path. + */ + private @Nullable String lastUnfinishedInvocationId(Session session) { + List events = session.events(); + String candidateId = null; + for (int i = events.size() - 1; i >= 0; i--) { + String id = events.get(i).invocationId(); + if (id != null && !id.isEmpty()) { + candidateId = id; + break; + } + } + if (candidateId == null) { + return null; + } + boolean hasCheckpoint = + events.stream() + .anyMatch(e -> e.actions().endOfAgent() || e.actions().agentState().isPresent()); + if (!hasCheckpoint) { + // Uncheckpointed sessions predate checkpoints and give no finished/paused signal, so resume + // only one still paused on an unanswered long-running call; else start a new invocation. + return hasUnansweredLongRunningCall(events, candidateId) ? candidateId : null; + } + // Finished when the agent a resume would run -- the active agent, not necessarily the root -- + // has emitted end-of-agent (mirrors runResumedAgent's guard). Keying on the root alone wedges + // after a transfer, where the root closes and later turns run the sub-agent. + String activeAgent = findAgentToRun(session, this.agent).name(); + for (Event event : events) { + if (candidateId.equals(event.invocationId()) + && activeAgent.equals(event.author()) + && event.actions().endOfAgent()) { + return null; + } + } + return candidateId; + } + + /** + * Returns whether {@code invocationId} holds a long-running function call (e.g. a HITL request or + * tool confirmation) with no matching function response yet -- the paused state a plain-text + * continuation may resume even when the session carries no resumability checkpoints. + */ + private static boolean hasUnansweredLongRunningCall(List events, String invocationId) { + Set answered = new HashSet<>(); + for (Event event : events) { + if (invocationId.equals(event.invocationId())) { + for (FunctionResponse response : event.functionResponses()) { + response.id().ifPresent(answered::add); + } + } + } + for (Event event : events) { + if (!invocationId.equals(event.invocationId())) { + continue; + } + Set longRunningIds = event.longRunningToolIds().orElse(ImmutableSet.of()); + for (FunctionCall call : event.functionCalls()) { + String id = call.id().orElse(null); + if (id != null && longRunningIds.contains(id) && !answered.contains(id)) { + return true; + } + } + } + return false; + } + + /** + * Returns the session event whose function call matches a function response id carried by {@code + * newMessage}, searching newest-first. Both the resumed invocation id and the branch of the + * appended function-response event are derived from it. + */ + private static Optional matchingFunctionCallEvent(Session session, Content newMessage) { + Set responseIds = new HashSet<>(); + newMessage + .parts() + .ifPresent( + parts -> + parts.forEach( + part -> + part.functionResponse() + .flatMap(FunctionResponse::id) + .ifPresent(responseIds::add))); + if (responseIds.isEmpty()) { + return Optional.empty(); + } + List events = session.events(); + for (int i = events.size() - 1; i >= 0; i--) { + Event event = events.get(i); + for (FunctionCall call : event.functionCalls()) { + if (call.id().filter(responseIds::contains).isPresent()) { + return Optional.of(event); + } + } + } + return Optional.empty(); + } + + /** Returns whether {@code message} carries any function response part. */ + private static boolean hasFunctionResponse(Content message) { + return message.parts().stream() + .flatMap(List::stream) + .anyMatch(part -> part.functionResponse().isPresent()); + } + + /** Returns whether the session holds at least one event belonging to {@code invocationId}. */ + private static boolean sessionHasEventsForInvocation(Session session, String invocationId) { + return session.events().stream().anyMatch(event -> invocationId.equals(event.invocationId())); + } + private void copySessionStates(Session source, Session target) { // TODO: remove this hack when deprecating all runAsync with Session. target.state().putAll(source.state()); @@ -717,7 +1128,6 @@ private InvocationContext.Builder newInvocationContextBuilder(Session session) { .artifactService(this.artifactService) .memoryService(this.memoryService) .pluginManager(this.pluginManager) - .agent(rootAgent) .session(session) .eventsCompactionConfig(this.eventsCompactionConfig) .contextCacheConfig(this.contextCacheConfig) @@ -809,10 +1219,7 @@ protected Flowable runLiveImpl( Span span = Span.current(); span.setStatus(StatusCode.ERROR, "Error in runLive Flowable execution"); span.recordException(throwable); - this.pluginManager - .runOnRunErrorCallback(invocationContext, throwable) - .onErrorComplete() - .subscribe(); + runOnRunError(invocationContext, throwable); }) .compose(Tracing.withContext(capturedContext)); }); @@ -845,19 +1252,26 @@ private boolean isResumable() { return resumabilityConfig != null && resumabilityConfig.isResumable(); } + /** + * @deprecated Back-compat only: see {@link + * com.google.adk.apps.ResumabilityConfig#isPlainTextContinuationAutoResume()}. + */ + @Deprecated + private boolean plainTextContinuationAutoResumes() { + return resumabilityConfig != null && resumabilityConfig.isPlainTextContinuationAutoResume(); + } + /** Returns the agent that should handle the next request based on session history. */ private BaseAgent findAgentToRun(Session session, BaseAgent rootAgent) { - // Route a function response to its call's author; when resumable, re-enter via the author's - // top-most SequentialAgent ancestor so the sequence can advance past it (else route straight to - // it, matching Python ADK v1 with resumability off). Temporary, event-based. + // Route a function response to its call's author; when resumable, re-enter via its + // top-most resume-aware workflow ancestor so the workflow can advance. Optional functionCallAuthor = Functions.findMatchingFunctionCallEvent(session.events()) .filter(event -> event.author() != null) .flatMap(event -> rootAgent.findAgent(event.author())); if (functionCallAuthor.isPresent()) { - return isResumable() - ? topmostSequentialAncestor(functionCallAuthor.get()) - : functionCallAuthor.get(); + BaseAgent author = functionCallAuthor.get(); + return isResumable() ? topmostResumableWorkflowAncestor(author) : author; } List events = new ArrayList<>(session.events()); @@ -872,6 +1286,13 @@ private BaseAgent findAgentToRun(Session session, BaseAgent rootAgent) { continue; } + // Skip resumability checkpoint markers (end_of_agent / agent_state): they carry no model + // turn, so a turn after a transfer resumes at the transferred-to sub-agent rather than the + // finished root. Such events exist only when resumable. + if (event.actions().endOfAgent() || event.actions().agentState().isPresent()) { + continue; + } + if (author.equals(rootAgent.name())) { return rootAgent; } @@ -891,15 +1312,15 @@ private BaseAgent findAgentToRun(Session session, BaseAgent rootAgent) { } /** - * Returns the top-most ancestor reachable from {@code agent} through {@link SequentialAgent} - * parents, or {@code agent} itself otherwise. Only SequentialAgent is resume-aware; other - * workflow agents are left to resume their paused sub-agent directly (via the function-call - * author). + * Returns the top-most ancestor reachable from {@code agent} through resume-aware workflow + * parents ({@link SequentialAgent} or {@link LoopAgent}), or {@code agent} itself otherwise, so a + * sub-agent resumed from a long-running pause re-enters the workflow that sequences it and the + * workflow can advance past it. Other agents resume their paused sub-agent directly. */ - private static BaseAgent topmostSequentialAncestor(BaseAgent agent) { + private static BaseAgent topmostResumableWorkflowAncestor(BaseAgent agent) { BaseAgent result = agent; BaseAgent parent = agent.parentAgent(); - while (parent instanceof SequentialAgent) { + while (parent instanceof SequentialAgent || parent instanceof LoopAgent) { result = parent; parent = parent.parentAgent(); } diff --git a/core/src/main/java/com/google/adk/sessions/SessionJsonConverter.java b/core/src/main/java/com/google/adk/sessions/SessionJsonConverter.java index adc84fbb6..8365117ef 100644 --- a/core/src/main/java/com/google/adk/sessions/SessionJsonConverter.java +++ b/core/src/main/java/com/google/adk/sessions/SessionJsonConverter.java @@ -114,6 +114,7 @@ static String convertEventToJson(Event event, boolean useIsoString) { putIfNotEmpty(actionsJson, "requestedAuthConfigs", actions.requestedAuthConfigs()); putIfNotEmpty( actionsJson, "requestedToolConfirmations", actions.requestedToolConfirmations()); + actions.agentState().ifPresent(v -> actionsJson.put("agentState", v)); eventJson.put("actions", actionsJson); } event.content().ifPresent(c -> eventJson.put("content", SessionUtils.encodeContent(c))); @@ -192,6 +193,10 @@ static Event fromApiEvent(Map apiEvent) { Optional.ofNullable(actionsMap.get("requestedToolConfirmations")) .map(SessionJsonConverter::asConcurrentMapOfToolConfirmations) .orElse(new ConcurrentHashMap<>())); + Object agentState = actionsMap.get("agentState"); + if (agentState != null) { + eventActionsBuilder.agentState((Map) agentState); + } } Event event = diff --git a/core/src/test/java/com/google/adk/agents/InvocationContextTest.java b/core/src/test/java/com/google/adk/agents/InvocationContextTest.java index e588a38ca..b87cac30f 100644 --- a/core/src/test/java/com/google/adk/agents/InvocationContextTest.java +++ b/core/src/test/java/com/google/adk/agents/InvocationContextTest.java @@ -20,15 +20,23 @@ import static org.junit.Assert.assertThrows; import static org.mockito.Mockito.mock; +import com.google.adk.apps.ResumabilityConfig; import com.google.adk.artifacts.BaseArtifactService; +import com.google.adk.events.Event; +import com.google.adk.events.EventActions; import com.google.adk.memory.BaseMemoryService; import com.google.adk.models.LlmCallsLimitExceededException; 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.genai.types.Content; +import com.google.genai.types.FunctionCall; +import com.google.genai.types.FunctionResponse; +import com.google.genai.types.Part; import java.util.HashMap; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; @@ -724,4 +732,452 @@ public void build_missingSessionService_throwsException() { IllegalStateException exception = assertThrows(IllegalStateException.class, builder::build); assertThat(exception).hasMessageThat().isEqualTo("Session service must be set."); } + + // ---- Resumability: runtime checkpoint state. ---- + + private InvocationContext resumableContext(Session eventSession, String invocationId) { + return InvocationContext.builder() + .sessionService(mockSessionService) + .artifactService(mockArtifactService) + .agent(mockAgent) + .session(eventSession) + .invocationId(invocationId) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build(); + } + + private static Event agentEvent( + String invocationId, String author, EventActions actions, Content content) { + Event.Builder builder = + Event.builder().id(Event.generateEventId()).invocationId(invocationId).author(author); + if (actions != null) { + builder.actions(actions); + } + if (content != null) { + builder.content(content); + } + return builder.build(); + } + + @Test + public void isResumable_configTrue_returnsTrue() { + InvocationContext context = resumableContext(session, "inv"); + assertThat(context.isResumable()).isTrue(); + } + + @Test + public void isResumable_configFalse_returnsFalse() { + InvocationContext context = + InvocationContext.builder() + .sessionService(mockSessionService) + .artifactService(mockArtifactService) + .agent(mockAgent) + .session(session) + .resumabilityConfig(ResumabilityConfig.builder().resumable(false).build()) + .build(); + assertThat(context.isResumable()).isFalse(); + } + + @Test + public void isResumable_nullConfig_returnsFalse() { + InvocationContext context = + InvocationContext.builder() + .sessionService(mockSessionService) + .artifactService(mockArtifactService) + .agent(mockAgent) + .session(session) + .build(); + assertThat(context.isResumable()).isFalse(); + } + + @Test + public void setAgentState_storesStateAndClearsEnd() { + InvocationContext context = resumableContext(session, "inv"); + + context.setAgentState("a", ImmutableMap.of("k", "v"), /* endOfAgent= */ false); + + assertThat(context.agentStates()).containsEntry("a", ImmutableMap.of("k", "v")); + assertThat(context.endOfAgents()).containsEntry("a", false); + } + + @Test + public void setAgentState_endOfAgent_marksEndedAndDropsState() { + InvocationContext context = resumableContext(session, "inv"); + context.setAgentState("a", ImmutableMap.of("k", "v"), /* endOfAgent= */ false); + + context.setAgentState("a", /* agentState= */ null, /* endOfAgent= */ true); + + assertThat(context.endOfAgents()).containsEntry("a", true); + assertThat(context.agentStates()).doesNotContainKey("a"); + } + + @Test + public void setAgentState_nullStateNotEnded_clearsBoth() { + InvocationContext context = resumableContext(session, "inv"); + context.setAgentState("a", ImmutableMap.of("k", "v"), /* endOfAgent= */ false); + + context.setAgentState("a", /* agentState= */ null, /* endOfAgent= */ false); + + assertThat(context.agentStates()).doesNotContainKey("a"); + assertThat(context.endOfAgents()).doesNotContainKey("a"); + } + + @Test + public void resetSubAgentStates_recursivelyClearsDescendants() { + BaseAgent grandChild = SequentialAgent.builder().name("gc").build(); + BaseAgent child1 = + SequentialAgent.builder().name("c1").subAgents(ImmutableList.of(grandChild)).build(); + BaseAgent child2 = SequentialAgent.builder().name("c2").build(); + BaseAgent parent = + SequentialAgent.builder().name("p").subAgents(ImmutableList.of(child1, child2)).build(); + InvocationContext context = + InvocationContext.builder() + .sessionService(mockSessionService) + .artifactService(mockArtifactService) + .agent(parent) + .session(session) + .invocationId("inv") + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build(); + context.setAgentState("c1", ImmutableMap.of("k", "v"), /* endOfAgent= */ false); + context.setAgentState("c2", ImmutableMap.of("k", "v"), /* endOfAgent= */ false); + context.setAgentState("gc", ImmutableMap.of("k", "v"), /* endOfAgent= */ false); + + context.resetSubAgentStates("p"); + + // Every descendant of p is cleared, including the grandchild reached recursively. + assertThat(context.agentStates()).doesNotContainKey("c1"); + assertThat(context.agentStates()).doesNotContainKey("c2"); + assertThat(context.agentStates()).doesNotContainKey("gc"); + } + + @Test + public void shouldPauseInvocation_resumableWithLongRunningCall_returnsTrue() { + InvocationContext context = resumableContext(session, "inv"); + Event event = + Event.builder() + .id("e1") + .invocationId("inv") + .author("a") + .content( + Content.fromParts( + Part.builder() + .functionCall(FunctionCall.builder().id("c1").name("tool").build()) + .build())) + .longRunningToolIds(ImmutableSet.of("c1")) + .build(); + + assertThat(context.shouldPauseInvocation(event)).isTrue(); + } + + @Test + public void shouldPauseInvocation_notResumable_returnsFalse() { + InvocationContext context = + InvocationContext.builder() + .sessionService(mockSessionService) + .artifactService(mockArtifactService) + .agent(mockAgent) + .session(session) + .invocationId("inv") + .build(); + Event event = + Event.builder() + .id("e1") + .invocationId("inv") + .author("a") + .content( + Content.fromParts( + Part.builder() + .functionCall(FunctionCall.builder().id("c1").name("tool").build()) + .build())) + .longRunningToolIds(ImmutableSet.of("c1")) + .build(); + + assertThat(context.shouldPauseInvocation(event)).isFalse(); + } + + @Test + public void shouldPauseInvocation_noLongRunningIds_returnsFalse() { + InvocationContext context = resumableContext(session, "inv"); + Event event = + Event.builder() + .id("e1") + .invocationId("inv") + .author("a") + .content( + Content.fromParts( + Part.builder() + .functionCall(FunctionCall.builder().id("c1").name("tool").build()) + .build())) + .build(); + + assertThat(context.shouldPauseInvocation(event)).isFalse(); + } + + @Test + public void shouldPauseInvocation_callIdNotInLongRunningSet_returnsFalse() { + InvocationContext context = resumableContext(session, "inv"); + Event event = + Event.builder() + .id("e1") + .invocationId("inv") + .author("a") + .content( + Content.fromParts( + Part.builder() + .functionCall(FunctionCall.builder().id("c1").name("tool").build()) + .build())) + .longRunningToolIds(ImmutableSet.of("other")) + .build(); + + assertThat(context.shouldPauseInvocation(event)).isFalse(); + } + + private static Event twoLongRunningCallsEvent() { + return Event.builder() + .id("m") + .invocationId("inv") + .author("root") + .content( + Content.fromParts( + Part.builder() + .functionCall(FunctionCall.builder().id("a").name("approve_a").build()) + .build(), + Part.builder() + .functionCall(FunctionCall.builder().id("b").name("approve_b").build()) + .build())) + .longRunningToolIds(ImmutableSet.of("a", "b")) + .build(); + } + + private static Event functionResponseEvent(String id, String name) { + return Event.builder() + .id("r-" + id) + .invocationId("inv") + .author("user") + .content( + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id(id) + .name(name) + .response(ImmutableMap.of("status", "done")) + .build()) + .build())) + .build(); + } + + @Test + public void hasUnansweredPausedCall_partiallyAnsweredParallelCalls_returnsTrue() { + Session eventSession = Session.builder("s").build(); + eventSession.events().add(twoLongRunningCallsEvent()); + eventSession.events().add(functionResponseEvent("a", "approve_a")); + InvocationContext context = resumableContext(eventSession, "inv"); + + assertThat(context.hasUnansweredPausedCall()).isTrue(); + } + + @Test + public void hasUnansweredPausedCall_allParallelCallsAnswered_returnsFalse() { + Session eventSession = Session.builder("s").build(); + eventSession.events().add(twoLongRunningCallsEvent()); + eventSession.events().add(functionResponseEvent("a", "approve_a")); + eventSession.events().add(functionResponseEvent("b", "approve_b")); + InvocationContext context = resumableContext(eventSession, "inv"); + + assertThat(context.hasUnansweredPausedCall()).isFalse(); + } + + @Test + public void hasUnansweredPausedCall_singleCallAnswered_returnsFalse() { + Session eventSession = Session.builder("s").build(); + eventSession + .events() + .add( + Event.builder() + .id("m") + .invocationId("inv") + .author("root") + .content( + Content.fromParts( + Part.builder() + .functionCall(FunctionCall.builder().id("c1").name("task").build()) + .build())) + .longRunningToolIds(ImmutableSet.of("c1")) + .build()); + eventSession.events().add(functionResponseEvent("c1", "task")); + InvocationContext context = resumableContext(eventSession, "inv"); + + assertThat(context.hasUnansweredPausedCall()).isFalse(); + } + + @Test + public void hasUnansweredPausedCall_noPausedCall_returnsFalse() { + Session eventSession = Session.builder("s").build(); + eventSession + .events() + .add(agentEvent("inv", "user", null, Content.fromParts(Part.fromText("hi")))); + eventSession + .events() + .add(agentEvent("inv", "root", null, Content.fromParts(Part.fromText("answer")))); + InvocationContext context = resumableContext(eventSession, "inv"); + + assertThat(context.hasUnansweredPausedCall()).isFalse(); + } + + @Test + public void events_filtersByInvocationAndBranch() { + Session eventSession = Session.builder("s").build(); + Event thisInv = agentEvent("inv", "a", null, Content.fromParts(Part.fromText("x"))); + Event otherInv = agentEvent("other", "a", null, Content.fromParts(Part.fromText("y"))); + Event branchB = + Event.builder() + .id("e3") + .invocationId("inv") + .author("a") + .branch("branchB") + .content(Content.fromParts(Part.fromText("z"))) + .build(); + eventSession.events().add(thisInv); + eventSession.events().add(otherInv); + eventSession.events().add(branchB); + InvocationContext context = resumableContext(eventSession, "inv"); + + assertThat(context.events(/* currentInvocation= */ true, /* currentBranch= */ false)) + .containsExactly(thisInv, branchB) + .inOrder(); + // A null-branch event is visible on any branch; the "branchB" event is filtered out. + assertThat(context.events(/* currentInvocation= */ true, /* currentBranch= */ true)) + .containsExactly(thisInv); + } + + @Test + public void populateInvocationAgentStates_notResumable_doesNothing() { + Session eventSession = Session.builder("s").build(); + eventSession + .events() + .add( + agentEvent( + "inv", + "a", + EventActions.builder().agentState(ImmutableMap.of("k", "v")).build(), + Content.fromParts(Part.fromText("x")))); + InvocationContext context = + InvocationContext.builder() + .sessionService(mockSessionService) + .artifactService(mockArtifactService) + .agent(mockAgent) + .session(eventSession) + .invocationId("inv") + .build(); + + context.populateInvocationAgentStates(); + + assertThat(context.agentStates()).isEmpty(); + assertThat(context.endOfAgents()).isEmpty(); + } + + @Test + public void populateInvocationAgentStates_endOfAgentEvent_marksEndedAndRemovesState() { + Session eventSession = Session.builder("s").build(); + eventSession + .events() + .add( + agentEvent( + "inv", + "a", + EventActions.builder().agentState(ImmutableMap.of("k", "v")).build(), + Content.fromParts(Part.fromText("x")))); + eventSession + .events() + .add(agentEvent("inv", "a", EventActions.builder().endOfAgent(true).build(), null)); + InvocationContext context = resumableContext(eventSession, "inv"); + + context.populateInvocationAgentStates(); + + assertThat(context.endOfAgents()).containsEntry("a", true); + assertThat(context.agentStates()).doesNotContainKey("a"); + } + + @Test + public void populateInvocationAgentStates_agentStateEvent_setsStateAndClearsEnd() { + Session eventSession = Session.builder("s").build(); + eventSession + .events() + .add( + agentEvent( + "inv", + "a", + EventActions.builder().agentState(ImmutableMap.of("k", "v")).build(), + Content.fromParts(Part.fromText("x")))); + InvocationContext context = resumableContext(eventSession, "inv"); + + context.populateInvocationAgentStates(); + + assertThat(context.agentStates()).containsEntry("a", ImmutableMap.of("k", "v")); + assertThat(context.endOfAgents()).containsEntry("a", false); + } + + @Test + public void populateInvocationAgentStates_agentStateAndEndOfAgent_endOfAgentWins() { + Session eventSession = Session.builder("s").build(); + eventSession + .events() + .add( + agentEvent( + "inv", + "a", + EventActions.builder() + .endOfAgent(true) + .agentState(ImmutableMap.of("k", "v")) + .build(), + Content.fromParts(Part.fromText("x")))); + InvocationContext context = resumableContext(eventSession, "inv"); + + context.populateInvocationAgentStates(); + + assertThat(context.endOfAgents()).containsEntry("a", true); + assertThat(context.agentStates()).doesNotContainKey("a"); + } + + @Test + public void populateInvocationAgentStates_newContentFromNonUserAuthor_initializesEmptyState() { + Session eventSession = Session.builder("s").build(); + eventSession + .events() + .add(agentEvent("inv", "a", null, Content.fromParts(Part.fromText("hello")))); + InvocationContext context = resumableContext(eventSession, "inv"); + + context.populateInvocationAgentStates(); + + assertThat(context.agentStates()).containsKey("a"); + assertThat(context.agentStates().get("a")).isEmpty(); + assertThat(context.endOfAgents()).containsEntry("a", false); + } + + @Test + public void populateInvocationAgentStates_userMessage_ignoredForDefaultState() { + Session eventSession = Session.builder("s").build(); + eventSession + .events() + .add(agentEvent("inv", "user", null, Content.fromParts(Part.fromText("hi")))); + InvocationContext context = resumableContext(eventSession, "inv"); + + context.populateInvocationAgentStates(); + + assertThat(context.agentStates()).isEmpty(); + } + + @Test + public void populateInvocationAgentStates_noContentNoState_ignored() { + Session eventSession = Session.builder("s").build(); + eventSession.events().add(agentEvent("inv", "a", null, null)); + InvocationContext context = resumableContext(eventSession, "inv"); + + context.populateInvocationAgentStates(); + + assertThat(context.agentStates()).isEmpty(); + assertThat(context.endOfAgents()).isEmpty(); + } } diff --git a/core/src/test/java/com/google/adk/agents/LlmAgentTest.java b/core/src/test/java/com/google/adk/agents/LlmAgentTest.java index 26843bb56..474352bae 100644 --- a/core/src/test/java/com/google/adk/agents/LlmAgentTest.java +++ b/core/src/test/java/com/google/adk/agents/LlmAgentTest.java @@ -19,13 +19,16 @@ import static com.google.adk.testing.TestUtils.assertEqualIgnoringFunctionIds; import static com.google.adk.testing.TestUtils.createInvocationContext; import static com.google.adk.testing.TestUtils.createLlmResponse; +import static com.google.adk.testing.TestUtils.createSubAgent; import static com.google.adk.testing.TestUtils.createTestAgent; import static com.google.adk.testing.TestUtils.createTestAgentBuilder; import static com.google.adk.testing.TestUtils.createTestLlm; import static com.google.adk.testing.TestUtils.createTextLlmResponse; +import static com.google.adk.testing.TestUtils.simplifyEvents; import static com.google.common.collect.Iterables.getOnlyElement; import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; import com.google.adk.agents.Callbacks.AfterModelCallback; import com.google.adk.agents.Callbacks.AfterToolCallback; @@ -33,7 +36,10 @@ import com.google.adk.agents.Callbacks.BeforeToolCallback; import com.google.adk.agents.Callbacks.OnModelErrorCallback; import com.google.adk.agents.Callbacks.OnToolErrorCallback; +import com.google.adk.apps.ResumabilityConfig; +import com.google.adk.artifacts.InMemoryArtifactService; import com.google.adk.events.Event; +import com.google.adk.events.EventActions; import com.google.adk.examples.Example; import com.google.adk.models.LlmRegistry; import com.google.adk.models.LlmRequest; @@ -49,8 +55,11 @@ import com.google.adk.tools.ExampleTool; 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 com.google.genai.types.Part; import com.google.genai.types.Schema; import com.google.genai.types.Type; @@ -632,4 +641,190 @@ public void run_withExampleTool_doesNotAddFunctionDeclarations() { var config = request.config().get(); assertThat(config.tools().isPresent()).isFalse(); } + + // ---- Resumability: resume into a transferred sub-agent (parity with ResumableLlmAgentTest). + // ---- + + private static InvocationContext resumableContextWithSeededEvent( + LlmAgent rootAgent, Event seededEvent) { + InMemorySessionService sessionService = new InMemorySessionService(); + Session session = sessionService.createSession("app", "user").blockingGet(); + var unused = sessionService.appendEvent(session, seededEvent).blockingGet(); + return InvocationContext.builder() + .sessionService(sessionService) + .artifactService(new InMemoryArtifactService()) + .invocationId("inv") + .agent(rootAgent) + .session(session) + .userContent(Content.fromParts(Part.fromText("hi"))) + .runConfig(RunConfig.builder().build()) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build(); + } + + @Test + public void runAsync_resumeFromTransferCall_runsTransferredSubAgent() { + LlmAgent sub = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("sub response"))) + .name("sub") + .build(); + TestLlm rootLlm = createTestLlm(createTextLlmResponse("root should not run")); + LlmAgent root = createTestAgentBuilder(rootLlm).name("root").subAgents(sub).build(); + Event transferEvent = + Event.builder() + .id("t1") + .invocationId("inv") + .author("root") + .actions(EventActions.builder().transferToAgent("sub").build()) + .content(Content.fromParts(Part.fromText("transferring"))) + .build(); + InvocationContext context = resumableContextWithSeededEvent(root, transferEvent); + context.setAgentState("root", ImmutableMap.of(), /* endOfAgent= */ false); + + List events = root.runAsync(context).toList().blockingGet(); + + // The transferred sub-agent runs; the root model is not re-invoked; root marks end-of-agent. + assertThat(simplifyEvents(events)).contains("sub: sub response"); + assertThat(rootLlm.getRequests()).isEmpty(); + assertThat( + events.stream() + .anyMatch(event -> event.author().equals("root") && event.actions().endOfAgent())) + .isTrue(); + // The end-of-agent is also recorded in the invocation state (not just the emitted event), so a + // later resume no-ops rather than re-running the completed root. + assertThat(context.endOfAgents()).containsEntry("root", true); + } + + @Test + public void runAsync_resumeUserResponseWithNoMatchingCall_throws() { + // The resume's last event is a user function response whose id matches no prior function call, + // so the resume-target resolution rejects it rather than silently continuing. + TestLlm rootLlm = createTestLlm(createTextLlmResponse("root should not run")); + LlmAgent root = createTestAgentBuilder(rootLlm).name("root").build(); + Event userResponse = + Event.builder() + .id("u1") + .invocationId("inv") + .author("user") + .content( + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("nomatch") + .name("tool") + .response(ImmutableMap.of("k", "v")) + .build()) + .build())) + .build(); + InvocationContext context = resumableContextWithSeededEvent(root, userResponse); + context.setAgentState("root", ImmutableMap.of(), /* endOfAgent= */ false); + + var resumed = root.runAsync(context).toList(); + assertThrows(IllegalArgumentException.class, resumed::blockingGet); + assertThat(rootLlm.getRequests()).isEmpty(); + } + + @Test + public void runAsync_resumeFromTransfer_subAgentRepauses_rootDoesNotEndOfAgent() { + // The transferred sub-agent pauses again on its own long-running call when resumed. + Event subPause = + Event.builder() + .id("p1") + .invocationId("inv") + .author("sub") + .content( + Content.builder() + .parts( + Part.builder() + .functionCall( + FunctionCall.builder() + .id("lro") + .name("waitTool") + .args(ImmutableMap.of()))) + .role("model") + .build()) + .longRunningToolIds(ImmutableSet.of("lro")) + .build(); + BaseAgent sub = createSubAgent("sub", subPause); + LlmAgent root = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("root should not run"))) + .name("root") + .subAgents(sub) + .build(); + Event transferEvent = + Event.builder() + .id("t1") + .invocationId("inv") + .author("root") + .actions(EventActions.builder().transferToAgent("sub").build()) + .content(Content.fromParts(Part.fromText("transferring"))) + .build(); + InvocationContext context = resumableContextWithSeededEvent(root, transferEvent); + context.setAgentState("root", ImmutableMap.of(), /* endOfAgent= */ false); + + List events = root.runAsync(context).toList().blockingGet(); + + // The transferred sub-agent re-pauses, so root must not mark end-of-agent (a later turn + // resumes that pause). + assertThat( + events.stream() + .anyMatch(event -> event.author().equals("root") && event.actions().endOfAgent())) + .isFalse(); + } + + @Test + public void runAsync_resumeNoTransfer_continuesRootAgent() { + TestLlm rootLlm = createTestLlm(createTextLlmResponse("root continues")); + LlmAgent root = createTestAgentBuilder(rootLlm).name("root").build(); + Event priorModelResponse = + Event.builder() + .id("m1") + .invocationId("inv") + .author("root") + .content(Content.fromParts(Part.fromText("earlier response"))) + .build(); + InvocationContext context = resumableContextWithSeededEvent(root, priorModelResponse); + context.setAgentState("root", ImmutableMap.of(), /* endOfAgent= */ false); + + List events = root.runAsync(context).toList().blockingGet(); + + // No transfer recorded: the root agent continues by invoking its model. + assertThat(simplifyEvents(events)).contains("root: root continues"); + assertThat(rootLlm.getRequests()).hasSize(1); + } + + @Test + public void runAsync_resumeFromTransferToPeer_runsTransferredPeerAgent() { + LlmAgent root = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("root"))) + .name("root") + .subAgents( + createTestAgentBuilder( + createTestLlm(createTextLlmResponse("agent A should not re-run"))) + .name("agent_a") + .build(), + createTestAgentBuilder(createTestLlm(createTextLlmResponse("agent B response"))) + .name("agent_b") + .build()) + .build(); + // agent_a transferred to its peer agent_b (not a descendant), then the invocation paused. + LlmAgent agentA = (LlmAgent) root.findAgent("agent_a").get(); + Event transferEvent = + Event.builder() + .id("t1") + .invocationId("inv") + .author("agent_a") + .actions(EventActions.builder().transferToAgent("agent_b").build()) + .content(Content.fromParts(Part.fromText("transferring to peer"))) + .build(); + InvocationContext context = resumableContextWithSeededEvent(agentA, transferEvent); + context.setAgentState("agent_a", ImmutableMap.of(), /* endOfAgent= */ false); + + List events = agentA.runAsync(context).toList().blockingGet(); + + // The transferred peer runs; agent_a does not re-run itself. + assertThat(simplifyEvents(events)).contains("agent_b: agent B response"); + assertThat(simplifyEvents(events)).doesNotContain("agent_a: agent A should not re-run"); + } } diff --git a/core/src/test/java/com/google/adk/agents/LoopAgentTest.java b/core/src/test/java/com/google/adk/agents/LoopAgentTest.java index b2d0778c6..7ef62f4ae 100644 --- a/core/src/test/java/com/google/adk/agents/LoopAgentTest.java +++ b/core/src/test/java/com/google/adk/agents/LoopAgentTest.java @@ -19,6 +19,7 @@ import static com.google.adk.testing.TestUtils.createEscalateEvent; import static com.google.adk.testing.TestUtils.createEvent; import static com.google.adk.testing.TestUtils.createInvocationContext; +import static com.google.adk.testing.TestUtils.createResumableInvocationContext; import static com.google.adk.testing.TestUtils.createSubAgent; import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.truth.Truth.assertThat; @@ -27,8 +28,11 @@ import com.google.adk.events.Event; import com.google.adk.testing.TestBaseAgent; 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.genai.types.Content; +import com.google.genai.types.FunctionCall; import com.google.genai.types.Part; import io.reactivex.rxjava3.core.Flowable; import io.reactivex.rxjava3.core.Maybe; @@ -240,4 +244,139 @@ public void runAsync_withEndInvocationInSubAgentCallback_stopsSubAgentButLoopCon assertThat(normalAgentRunCount.get()).isEqualTo(3); assertThat(subAgent2RunCount.get()).isEqualTo(1); } + + // ---- Resumability: durable checkpoint resume. ---- + + @Test + public void runAsync_resumable_emitsEndOfAgent() { + TestBaseAgent subAgent = + createSubAgent("sub", createEvent("e").toBuilder().author("sub").build()); + LoopAgent loopAgent = + LoopAgent.builder().name("loop").subAgents(subAgent).maxIterations(1).build(); + InvocationContext context = createResumableInvocationContext(loopAgent); + + List events = loopAgent.runAsync(context).toList().blockingGet(); + + assertThat( + events.stream() + .anyMatch(event -> event.author().equals("loop") && event.actions().endOfAgent())) + .isTrue(); + } + + @Test + public void runAsync_notResumable_doesNotEmitEndOfAgent() { + TestBaseAgent subAgent = + createSubAgent("sub", createEvent("e").toBuilder().author("sub").build()); + LoopAgent loopAgent = + LoopAgent.builder().name("loop").subAgents(subAgent).maxIterations(1).build(); + InvocationContext context = createInvocationContext(loopAgent); + + List events = loopAgent.runAsync(context).toList().blockingGet(); + + assertThat(events.stream().anyMatch(event -> event.actions().endOfAgent())).isFalse(); + } + + @Test + public void runAsync_resumingFromMiddle_restoresIterationAndAgent() { + TestBaseAgent agent1 = + createSubAgent( + "agent1", () -> Flowable.just(createEvent("a1").toBuilder().author("agent1").build())); + TestBaseAgent agent2 = + createSubAgent( + "agent2", () -> Flowable.just(createEvent("a2").toBuilder().author("agent2").build())); + LoopAgent loopAgent = + LoopAgent.builder().name("loop").subAgents(agent1, agent2).maxIterations(3).build(); + InvocationContext context = createResumableInvocationContext(loopAgent); + // Resume mid-loop: iteration 1, at agent2. + context.setAgentState( + "loop", + ImmutableMap.of("current_sub_agent", "agent2", "times_looped", 1), + /* endOfAgent= */ false); + + List events = loopAgent.runAsync(context).toList().blockingGet(); + + // The first sub-agent to run is agent2 (agent1 is skipped in the resumed iteration). + String firstSubAgentAuthor = + events.stream() + .filter(event -> event.content().isPresent()) + .map(Event::author) + .filter(author -> author.equals("agent1") || author.equals("agent2")) + .findFirst() + .orElse(null); + assertThat(firstSubAgentAuthor).isEqualTo("agent2"); + // The loop still finishes (reaches maxIterations) and marks end-of-agent. + assertThat( + events.stream() + .anyMatch(event -> event.author().equals("loop") && event.actions().endOfAgent())) + .isTrue(); + } + + @Test + public void runAsync_resumingAtIterationCap_endsWithoutRunningSubAgent() { + TestBaseAgent subAgent = + createSubAgent("sub", createEvent("e").toBuilder().author("sub").build()); + LoopAgent loopAgent = + LoopAgent.builder().name("loop").subAgents(subAgent).maxIterations(2).build(); + InvocationContext context = createResumableInvocationContext(loopAgent); + // Resume with the loop already at the iteration cap (2 of 2 done). + context.setAgentState( + "loop", + ImmutableMap.of("current_sub_agent", "sub", "times_looped", 2), + /* endOfAgent= */ false); + + List events = loopAgent.runAsync(context).toList().blockingGet(); + + // Ends immediately: end-of-agent is emitted and no sub-agent runs. + assertThat( + events.stream() + .anyMatch(event -> event.author().equals("loop") && event.actions().endOfAgent())) + .isTrue(); + assertThat(events.stream().anyMatch(event -> event.author().equals("sub"))).isFalse(); + } + + @Test + public void runAsync_resumable_belowIterationCap_loopsExactlyMaxIterations() { + // A fresh resumable run below the cap must loop exactly maxIterations times, guarding the cap + // comparison timesLooped >= maxIterations; complements runAsync_resumingAtIterationCap_*. + TestBaseAgent subAgent = + createSubAgent( + "sub", () -> Flowable.just(createEvent("e").toBuilder().author("sub").build())); + LoopAgent loopAgent = + LoopAgent.builder().name("loop").subAgents(subAgent).maxIterations(2).build(); + InvocationContext context = createResumableInvocationContext(loopAgent); + + List events = loopAgent.runAsync(context).toList().blockingGet(); + + assertThat(events.stream().filter(event -> event.author().equals("sub")).count()).isEqualTo(2L); + assertThat( + events.stream() + .anyMatch(event -> event.author().equals("loop") && event.actions().endOfAgent())) + .isTrue(); + } + + @Test + public void runAsync_resumable_pausesOnLongRunningCall_doesNotEmitEndOfAgent() { + Event longRunningCall = + Event.builder() + .id("lro") + .invocationId("invocationId") + .author("sub") + .content( + Content.fromParts( + Part.builder() + .functionCall(FunctionCall.builder().id("c1").name("tool").build()) + .build())) + .longRunningToolIds(ImmutableSet.of("c1")) + .build(); + TestBaseAgent subAgent = createSubAgent("sub", longRunningCall); + LoopAgent loopAgent = + LoopAgent.builder().name("loop").subAgents(subAgent).maxIterations(5).build(); + InvocationContext context = createResumableInvocationContext(loopAgent); + + List events = loopAgent.runAsync(context).toList().blockingGet(); + + // Paused on the long-running call: no end-of-agent, and the loop did not run 5 iterations. + assertThat(events.stream().anyMatch(event -> event.actions().endOfAgent())).isFalse(); + assertThat(events.stream().filter(event -> event.author().equals("sub")).count()).isEqualTo(1L); + } } diff --git a/core/src/test/java/com/google/adk/agents/ParallelAgentTest.java b/core/src/test/java/com/google/adk/agents/ParallelAgentTest.java index e51240c45..f4336f419 100644 --- a/core/src/test/java/com/google/adk/agents/ParallelAgentTest.java +++ b/core/src/test/java/com/google/adk/agents/ParallelAgentTest.java @@ -16,12 +16,22 @@ package com.google.adk.agents; +import static com.google.adk.testing.TestUtils.createEscalateEvent; +import static com.google.adk.testing.TestUtils.createFunctionCallLlmResponse; import static com.google.adk.testing.TestUtils.createInvocationContext; +import static com.google.adk.testing.TestUtils.createResumableInvocationContext; +import static com.google.adk.testing.TestUtils.createSubAgent; +import static com.google.adk.testing.TestUtils.createTestAgentBuilder; +import static com.google.adk.testing.TestUtils.createTestLlm; +import static com.google.adk.testing.TestUtils.createTextLlmResponse; +import static com.google.adk.testing.TestUtils.simplifyEvents; import static com.google.common.truth.Truth.assertThat; import static java.util.concurrent.TimeUnit.MILLISECONDS; import com.google.adk.events.Event; +import com.google.adk.tools.FunctionTool; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; import com.google.genai.types.Content; import com.google.genai.types.Part; import io.reactivex.rxjava3.core.Flowable; @@ -30,6 +40,7 @@ import io.reactivex.rxjava3.schedulers.TestScheduler; import io.reactivex.rxjava3.subscribers.TestSubscriber; import java.util.List; +import org.jspecify.annotations.Nullable; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -194,4 +205,212 @@ public void runAsync_withTestScheduler_usesVirtualTime() { testSubscriber.assertValueCount(1); testSubscriber.assertComplete(); } + + // ---- Resumability: checkpoint / skip-completed / pause. ---- + + /** Tools for the resumability tests. */ + public static final class Tools { + private Tools() {} + + // A long-running tool awaiting an external result returns nothing yet, so a branch pauses. + @SuppressWarnings("unused") // Invoked reflectively by FunctionTool. + public static @Nullable ImmutableMap waitForApproval(String reason) { + return null; + } + } + + @Test + public void runAsync_resumable_allSubAgentsEnd_emitsEndOfAgent() { + LlmAgent sub1 = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("sub1 done"))) + .name("sub1") + .build(); + LlmAgent sub2 = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("sub2 done"))) + .name("sub2") + .build(); + ParallelAgent parallelAgent = + ParallelAgent.builder().name("parallel").subAgents(sub1, sub2).build(); + InvocationContext context = createResumableInvocationContext(parallelAgent); + + List events = parallelAgent.runAsync(context).toList().blockingGet(); + + assertThat( + events.stream() + .anyMatch( + event -> event.author().equals("parallel") && event.actions().endOfAgent())) + .isTrue(); + } + + @Test + public void runAsync_resumable_sequentialSubAgentEnds_emitsEndOfAgent() { + LlmAgent leaf = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("leaf done"))) + .name("leaf") + .build(); + SequentialAgent sequential = + SequentialAgent.builder().name("sequential").subAgents(leaf).build(); + ParallelAgent parallelAgent = + ParallelAgent.builder().name("parallel").subAgents(sequential).build(); + InvocationContext context = createResumableInvocationContext(parallelAgent); + + List events = parallelAgent.runAsync(context).toList().blockingGet(); + + // A SequentialAgent child records end-of-agent on completion, so the parallel agent ends too. + assertThat( + events.stream() + .anyMatch( + event -> event.author().equals("parallel") && event.actions().endOfAgent())) + .isTrue(); + } + + @Test + public void runAsync_resumable_loopSubAgentEnds_emitsEndOfAgent() { + LlmAgent leaf = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("leaf done"))) + .name("leaf") + .build(); + LoopAgent loop = LoopAgent.builder().name("loop").subAgents(leaf).maxIterations(1).build(); + ParallelAgent parallelAgent = ParallelAgent.builder().name("parallel").subAgents(loop).build(); + InvocationContext context = createResumableInvocationContext(parallelAgent); + + List events = parallelAgent.runAsync(context).toList().blockingGet(); + + // A LoopAgent child records end-of-agent on completion, so the parallel agent ends too. + assertThat( + events.stream() + .anyMatch( + event -> event.author().equals("parallel") && event.actions().endOfAgent())) + .isTrue(); + } + + @Test + public void runAsync_resumable_subAgentEscalates_emitsEndOfAgent() { + BaseAgent escalating = createSubAgent("escalating", createEscalateEvent("esc")); + ParallelAgent parallelAgent = + ParallelAgent.builder().name("parallel").subAgents(escalating).build(); + InvocationContext context = createResumableInvocationContext(parallelAgent); + + List events = parallelAgent.runAsync(context).toList().blockingGet(); + + // A sub-agent escalation ends the parallel agent even though no branch recorded end-of-agent. + assertThat( + events.stream() + .anyMatch( + event -> event.author().equals("parallel") && event.actions().endOfAgent())) + .isTrue(); + } + + @Test + public void runAsync_notResumable_doesNotEmitEndOfAgent() { + LlmAgent sub1 = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("sub1 done"))) + .name("sub1") + .build(); + ParallelAgent parallelAgent = ParallelAgent.builder().name("parallel").subAgents(sub1).build(); + InvocationContext context = createInvocationContext(parallelAgent); + + List events = parallelAgent.runAsync(context).toList().blockingGet(); + + assertThat(events.stream().anyMatch(event -> event.actions().endOfAgent())).isFalse(); + } + + @Test + public void runAsync_resumable_oneBranchPauses_doesNotEmitEndOfAgent() { + LlmAgent pausing = + createTestAgentBuilder( + createTestLlm( + createFunctionCallLlmResponse( + "c1", "waitForApproval", ImmutableMap.of("reason", "x")))) + .name("pausing") + .tools( + FunctionTool.create( + Tools.class, + "waitForApproval", + /* requireConfirmation= */ false, + /* isLongRunning= */ true)) + .build(); + LlmAgent completing = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("done"))) + .name("completing") + .build(); + ParallelAgent parallelAgent = + ParallelAgent.builder().name("parallel").subAgents(pausing, completing).build(); + InvocationContext context = createResumableInvocationContext(parallelAgent); + + List events = parallelAgent.runAsync(context).toList().blockingGet(); + + // One branch paused: the parallel agent does not end. + assertThat( + events.stream() + .anyMatch( + event -> event.author().equals("parallel") && event.actions().endOfAgent())) + .isFalse(); + } + + @Test + public void runAsync_resumable_firstRun_emitsInitialStateCheckpoint() { + LlmAgent sub1 = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("sub1 done"))) + .name("sub1") + .build(); + ParallelAgent parallelAgent = ParallelAgent.builder().name("parallel").subAgents(sub1).build(); + InvocationContext context = createResumableInvocationContext(parallelAgent); + // No prior state for "parallel": this is a first run, so an initial checkpoint is recorded. + + List events = parallelAgent.runAsync(context).toList().blockingGet(); + + assertThat( + events.stream() + .anyMatch( + event -> + event.author().equals("parallel") + && event.actions().agentState().isPresent())) + .isTrue(); + } + + @Test + public void runAsync_resumable_resume_doesNotReEmitInitialStateCheckpoint() { + LlmAgent sub1 = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("sub1 done"))) + .name("sub1") + .build(); + ParallelAgent parallelAgent = ParallelAgent.builder().name("parallel").subAgents(sub1).build(); + InvocationContext context = createResumableInvocationContext(parallelAgent); + // The parallel agent already checkpointed that it started in a prior run. + context.setAgentState("parallel", ImmutableMap.of(), /* endOfAgent= */ false); + + List events = parallelAgent.runAsync(context).toList().blockingGet(); + + // No new initial checkpoint is emitted on resume (only the final end-of-agent checkpoint). + assertThat( + events.stream() + .anyMatch( + event -> + event.author().equals("parallel") + && event.actions().agentState().isPresent())) + .isFalse(); + } + + @Test + public void runAsync_resumable_skipsCompletedBranchesOnResume() { + LlmAgent sub1 = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("sub1 done"))) + .name("sub1") + .build(); + LlmAgent sub2 = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("sub2 done"))) + .name("sub2") + .build(); + ParallelAgent parallelAgent = + ParallelAgent.builder().name("parallel").subAgents(sub1, sub2).build(); + InvocationContext context = createResumableInvocationContext(parallelAgent); + // sub1 already finished in a prior run. + context.setAgentState("sub1", /* agentState= */ null, /* endOfAgent= */ true); + + List events = parallelAgent.runAsync(context).toList().blockingGet(); + + assertThat(simplifyEvents(events)).doesNotContain("sub1: sub1 done"); + assertThat(simplifyEvents(events)).contains("sub2: sub2 done"); + } } diff --git a/core/src/test/java/com/google/adk/agents/SequentialAgentTest.java b/core/src/test/java/com/google/adk/agents/SequentialAgentTest.java index 6bbd9e55b..3ee551f20 100644 --- a/core/src/test/java/com/google/adk/agents/SequentialAgentTest.java +++ b/core/src/test/java/com/google/adk/agents/SequentialAgentTest.java @@ -19,12 +19,18 @@ import static com.google.adk.testing.TestUtils.createEvent; import static com.google.adk.testing.TestUtils.createInvocationContext; import static com.google.adk.testing.TestUtils.createLlmResponse; +import static com.google.adk.testing.TestUtils.createResumableInvocationContext; import static com.google.adk.testing.TestUtils.createSubAgent; import static com.google.adk.testing.TestUtils.createTestAgent; +import static com.google.adk.testing.TestUtils.createTestAgentBuilder; import static com.google.adk.testing.TestUtils.createTestLlm; +import static com.google.adk.testing.TestUtils.createTextLlmResponse; +import static com.google.adk.testing.TestUtils.simplifyEvents; import static com.google.common.collect.Iterables.getOnlyElement; import static com.google.common.truth.Truth.assertThat; +import com.google.adk.apps.ResumabilityConfig; +import com.google.adk.artifacts.InMemoryArtifactService; import com.google.adk.events.Event; import com.google.adk.sessions.InMemorySessionService; import com.google.adk.sessions.Session; @@ -32,6 +38,7 @@ import com.google.adk.testing.TestLlm; 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.FunctionCall; import com.google.genai.types.FunctionResponse; @@ -284,4 +291,137 @@ private static InvocationContext contextResumingCall(BaseAgent rootAgent, String var unusedResponse = sessionService.appendEvent(session, responseEvent).blockingGet(); return createInvocationContext(rootAgent, sessionService, session); } + + // ---- Resumability: durable checkpoint resume. ---- + + @Test + public void runAsync_resumingFromMiddle_startsFromCorrectAgent() { + LlmAgent agent1 = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("a1 done"))) + .name("agent1") + .build(); + LlmAgent agent2 = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("a2 done"))) + .name("agent2") + .build(); + LlmAgent agent3 = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("a3 done"))) + .name("agent3") + .build(); + SequentialAgent sequentialAgent = + SequentialAgent.builder().name("seq").subAgents(agent1, agent2, agent3).build(); + InvocationContext context = createResumableInvocationContext(sequentialAgent); + // Seed the checkpoint: resume at agent2. + context.setAgentState( + "seq", ImmutableMap.of("current_sub_agent", "agent2"), /* endOfAgent= */ false); + + List events = sequentialAgent.runAsync(context).toList().blockingGet(); + + assertThat(simplifyEvents(events)).doesNotContain("agent1: a1 done"); + assertThat(simplifyEvents(events)).contains("agent2: a2 done"); + assertThat(simplifyEvents(events)).contains("agent3: a3 done"); + } + + @Test + public void runAsync_resumable_emitsEndOfAgent() { + LlmAgent agent1 = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("done"))).name("agent1").build(); + SequentialAgent sequentialAgent = + SequentialAgent.builder().name("seq").subAgents(agent1).build(); + InvocationContext context = createResumableInvocationContext(sequentialAgent); + + List events = sequentialAgent.runAsync(context).toList().blockingGet(); + + assertThat( + events.stream() + .anyMatch(event -> event.author().equals("seq") && event.actions().endOfAgent())) + .isTrue(); + } + + @Test + public void runAsync_notResumable_doesNotEmitEndOfAgent() { + LlmAgent agent1 = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("done"))).name("agent1").build(); + SequentialAgent sequentialAgent = + SequentialAgent.builder().name("seq").subAgents(agent1).build(); + InvocationContext context = createInvocationContext(sequentialAgent); + + List events = sequentialAgent.runAsync(context).toList().blockingGet(); + + assertThat(events.stream().anyMatch(event -> event.actions().endOfAgent())).isFalse(); + } + + // Back-compat: a session paused before checkpoints existed must still fast-forward past + // completed sub-agents (reconstructed from history) rather than re-running them. + @Test + public void runAsync_resumeLegacySessionWithoutCheckpoints_doesNotRerunCompletedSubAgents() { + TestBaseAgent agent1 = + createSubAgent("agent1", createEvent("a1").toBuilder().author("agent1").build()); + TestBaseAgent agent2 = + createSubAgent("agent2", createEvent("a2").toBuilder().author("agent2").build()); + TestBaseAgent agent3 = + createSubAgent("agent3", createEvent("a3").toBuilder().author("agent3").build()); + SequentialAgent sequentialAgent = + SequentialAgent.builder().name("seq").subAgents(agent1, agent2, agent3).build(); + + // Simulate a legacy paused session: agent2 issued a long-running call (no checkpoint events), + // and the user has now supplied the awaited function response. + InMemorySessionService sessionService = new InMemorySessionService(); + Session session = sessionService.createSession("app", "user").blockingGet(); + var unusedCall = + sessionService + .appendEvent( + session, + Event.builder() + .id("fc") + .invocationId("inv") + .author("agent2") + .content( + Content.fromParts( + Part.builder() + .functionCall(FunctionCall.builder().id("c1").name("tool").build()) + .build())) + .longRunningToolIds(ImmutableSet.of("c1")) + .build()) + .blockingGet(); + var unusedResponse = + sessionService + .appendEvent( + session, + Event.builder() + .id("fr") + .invocationId("inv") + .author("user") + .content( + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("c1") + .name("tool") + .response(ImmutableMap.of()) + .build()) + .build())) + .build()) + .blockingGet(); + InvocationContext context = + InvocationContext.builder() + .sessionService(sessionService) + .artifactService(new InMemoryArtifactService()) + .invocationId("inv") + .agent(sequentialAgent) + .session(session) + .userContent(Content.fromParts(Part.fromText("resume"))) + .runConfig(RunConfig.builder().build()) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build(); + // No populate/checkpoint state seeded -- this is the legacy case. + + var unused = sequentialAgent.runAsync(context).toList().blockingGet(); + + // agent1 (already completed before the pause) must not re-run; agent2 and agent3 do. + assertThat(agent1.getInvocationCount()).isEqualTo(0); + assertThat(agent2.getInvocationCount()).isEqualTo(1); + assertThat(agent3.getInvocationCount()).isEqualTo(1); + } } diff --git a/core/src/test/java/com/google/adk/flows/llmflows/FunctionsTest.java b/core/src/test/java/com/google/adk/flows/llmflows/FunctionsTest.java index 8e8555114..d0d04bdc3 100644 --- a/core/src/test/java/com/google/adk/flows/llmflows/FunctionsTest.java +++ b/core/src/test/java/com/google/adk/flows/llmflows/FunctionsTest.java @@ -399,6 +399,28 @@ public void getAskUserConfirmationFunctionCalls_eventWithConfirmationFunctionCal assertThat(result).containsExactly(confirmationCall1, confirmationCall2); } + @Test + public void hasPendingLongRunningCall_singleEventWithFunctionResponse_returnsFalse() { + // A single event cannot hold both a paused call and its response, so nothing is pending (and + // the pending-call lookup must not read a nonexistent prior event). + assertThat(Functions.hasPendingLongRunningCall(ImmutableList.of(functionResponseEvent("c1")))) + .isFalse(); + } + + @Test + public void hasPendingLongRunningCall_sequentialAnsweredCalls_returnsFalse() { + // Two long-running calls, each answered by the immediately following response: the pending-call + // lookup must inspect the second-to-last event (the latest call), not some other index. + assertThat( + Functions.hasPendingLongRunningCall( + ImmutableList.of( + longRunningCallEvent("c1"), + functionResponseEvent("c1"), + longRunningCallEvent("c2"), + functionResponseEvent("c2")))) + .isFalse(); + } + // Default ToolExecutionMode.NONE behaves like PARALLEL: blocking tools still execute serially // on the caller thread (no worker scheduler is used), preserving the historical default. @Test @@ -568,10 +590,44 @@ public void hasPendingLongRunningCall_emptyList_returnsFalse() { assertThat(Functions.hasPendingLongRunningCall(ImmutableList.of())).isFalse(); } + @Test + public void hasPendingLongRunningCall_list_responseResolvesCall_returnsFalse() { + // The trailing function response resolves the pending long-running call, so the flow continues. + ImmutableList events = + ImmutableList.of(longRunningCallEvent("call1"), functionResponseEvent("call1")); + assertThat(Functions.hasPendingLongRunningCall(events)).isFalse(); + } + + @Test + public void hasPendingLongRunningCall_list_responseForDifferentCall_returnsTrue() { + // The response does not resolve the pending call, so the long-running call still pauses. + ImmutableList events = + ImmutableList.of(longRunningCallEvent("call1"), functionResponseEvent("other")); + assertThat(Functions.hasPendingLongRunningCall(events)).isTrue(); + } + private static Event longRunningCallEvent(String callId) { return functionCallEvent(callId, callId); } + private static Event functionResponseEvent(String callId) { + return Event.builder() + .id("response_" + callId) + .invocationId("invocation1") + .author("user") + .content( + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id(callId) + .name("tool") + .response(ImmutableMap.of()) + .build()) + .build())) + .build(); + } + // Event with a function call; longRunningId, when non-null, is marked long-running. private static Event functionCallEvent(String callId, String longRunningId) { Event.Builder builder = diff --git a/core/src/test/java/com/google/adk/runner/RunnerTest.java b/core/src/test/java/com/google/adk/runner/RunnerTest.java index 3870d3461..bda21e1a4 100644 --- a/core/src/test/java/com/google/adk/runner/RunnerTest.java +++ b/core/src/test/java/com/google/adk/runner/RunnerTest.java @@ -23,16 +23,19 @@ import static com.google.adk.testing.TestUtils.createTestLlm; import static com.google.adk.testing.TestUtils.createTextLlmResponse; import static com.google.adk.testing.TestUtils.simplifyEvents; +import static com.google.adk.testing.TestUtils.simplifyResumableEvents; import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.truth.Truth.assertThat; import static java.nio.charset.StandardCharsets.UTF_8; import static java.util.Arrays.stream; import static java.util.concurrent.TimeUnit.MILLISECONDS; import static java.util.concurrent.TimeUnit.SECONDS; +import static org.junit.Assert.assertThrows; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.CALLS_REAL_METHODS; +import static org.mockito.Mockito.atLeastOnce; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.times; @@ -54,6 +57,7 @@ import com.google.adk.artifacts.BaseArtifactService; import com.google.adk.artifacts.InMemoryArtifactService; import com.google.adk.events.Event; +import com.google.adk.events.EventActions; import com.google.adk.flows.llmflows.Functions; import com.google.adk.models.LlmRequest; import com.google.adk.models.LlmResponse; @@ -67,6 +71,7 @@ import com.google.adk.sessions.SessionKey; import com.google.adk.summarizer.EventsCompactionConfig; import com.google.adk.telemetry.Tracing; +import com.google.adk.testing.TestBaseAgent; import com.google.adk.testing.TestLlm; import com.google.adk.testing.TestUtils; import com.google.adk.testing.TestUtils.EchoTool; @@ -77,6 +82,7 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.Iterables; +import com.google.common.collect.Streams; import com.google.genai.types.Content; import com.google.genai.types.FunctionCall; import com.google.genai.types.FunctionDeclaration; @@ -2428,12 +2434,17 @@ public void runAsync_withToolConfirmation_inSequentialAgent_runsLaterSubAgentsAf .toList() .blockingGet(); - // Turn 2: B resumes and executes the tool, then C runs. A is not re-run. - assertThat(simplifyEvents(eventsAfterConfirmation)) + // Turn 2: B resumes and executes the tool, then C runs (A is not re-run), with per-agent and + // workflow checkpoints. + assertThat(simplifyResumableEvents(eventsAfterConfirmation)) .containsExactly( "b_agent: FunctionResponse(name=echoTool, response={message=hello})", "b_agent: Response after user confirmed.", - "c_agent: agent C done") + "b_agent: end_of_agent", + "workflow_agent: agent_state={current_sub_agent=c_agent}", + "c_agent: agent C done", + "c_agent: end_of_agent", + "workflow_agent: end_of_agent") .inOrder(); } @@ -2446,104 +2457,1646 @@ public void runAsync_withLongRunningCall_inSequentialAgent_runsLaterSubAgentsAft createTestAgentBuilder(createTestLlm(createTextLlmResponse("agent A done"))) .name("a_agent") .build(); - // With resumability on, B pauses right after the long-running call (no extra model call), so a - // single follow-up response covers the resume. + // With resumability on, B pauses right after the no-result long-running call (no extra model + // call), so a single follow-up response covers the resume. TestLlm bTestLlm = createTestLlm( createFunctionCallLlmResponse( - "lro_call_id", "echoTool", ImmutableMap.of("message", "hello")), - createTextLlmResponse("agent B resumed")); - LlmAgent agentB = - createTestAgentBuilder(bTestLlm) - .name("b_agent") + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hello")), + createTextLlmResponse("agent B resumed")); + LlmAgent agentB = + createTestAgentBuilder(bTestLlm) + .name("b_agent") + .tools( + FunctionTool.create( + Tools.class, + "pendingTool", + /* requireConfirmation= */ false, + /* isLongRunning= */ true)) + .build(); + LlmAgent agentC = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("agent C done"))) + .name("c_agent") + .build(); + SequentialAgent workflowAgent = + SequentialAgent.builder() + .name("workflow_agent") + .subAgents(ImmutableList.of(agentA, agentB, agentC)) + .build(); + Runner runner = + Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(workflowAgent) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build()) + .build(); + Session session = runner.sessionService().createSession("test", "user").blockingGet(); + + List eventsBeforeResume = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) + .toList() + .blockingGet(); + + // Turn 1: A runs, B issues the long-running call and pauses; C must not run yet. B must not + // make + // a further model call after the pending call. + assertThat(simplifyEvents(eventsBeforeResume)).contains("a_agent: agent A done"); + assertThat(simplifyEvents(eventsBeforeResume)).doesNotContain("b_agent: agent B resumed"); + assertThat(simplifyEvents(eventsBeforeResume)).doesNotContain("c_agent: agent C done"); + + List eventsAfterResume = + runner + .runAsync( + "user", + session.id(), + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("lro_call_id") + .name("pendingTool") + .response(ImmutableMap.of("message", "hello"))) + .build())) + .toList() + .blockingGet(); + + // Turn 2: B resumes from the long-running response and C runs (A is not re-run), with per-agent + // and workflow checkpoints. + assertThat(simplifyResumableEvents(eventsAfterResume)) + .containsExactly( + "b_agent: agent B resumed", + "b_agent: end_of_agent", + "workflow_agent: agent_state={current_sub_agent=c_agent}", + "c_agent: agent C done", + "c_agent: end_of_agent", + "workflow_agent: end_of_agent") + .inOrder(); + } + + // A resumable LoopAgent(w1, w2) paused on w1's long-running call resumes w1 and then advances the + // loop to w2 and closes it, rather than resuming only the paused sub-agent -- the loop advances + // like a SequentialAgent. + @Test + @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). + public void runAsync_withLongRunningCall_inLoopAgent_runsRemainingSubAgentsAfterResume() { + TestLlm w1TestLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hello")), + createTextLlmResponse("w1 resumed")); + LlmAgent w1 = + createTestAgentBuilder(w1TestLlm) + .name("w1_agent") + .tools( + FunctionTool.create( + Tools.class, + "pendingTool", + /* requireConfirmation= */ false, + /* isLongRunning= */ true)) + .build(); + LlmAgent w2 = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("w2 done"))) + .name("w2_agent") + .build(); + LoopAgent workflowAgent = + LoopAgent.builder() + .name("loop_agent") + .subAgents(ImmutableList.of(w1, w2)) + .maxIterations(1) + .build(); + Runner runner = + Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(workflowAgent) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build()) + .build(); + Session session = runner.sessionService().createSession("test", "user").blockingGet(); + + List eventsBeforeResume = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) + .toList() + .blockingGet(); + + // Turn 1: w1 issues the long-running call and pauses; w2 must not run yet. + assertThat(simplifyEvents(eventsBeforeResume)).doesNotContain("w1_agent: w1 resumed"); + assertThat(simplifyEvents(eventsBeforeResume)).doesNotContain("w2_agent: w2 done"); + + List eventsAfterResume = + runner + .runAsync( + "user", + session.id(), + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("lro_call_id") + .name("pendingTool") + .response(ImmutableMap.of("message", "hello"))) + .build())) + .toList() + .blockingGet(); + + // Turn 2: w1 resumes and the loop advances to w2 and closes, with per-agent and loop + // checkpoints (w1 is not re-run from the start of the iteration). + assertThat(simplifyResumableEvents(eventsAfterResume)) + .containsExactly( + "w1_agent: w1 resumed", + "w1_agent: end_of_agent", + "loop_agent: agent_state={current_sub_agent=w2_agent, times_looped=0}", + "w2_agent: w2 done", + "w2_agent: end_of_agent", + "loop_agent: end_of_agent") + .inOrder(); + } + + // A resumable plain-agent transfer closes the root (end_of_agent) the moment it transfers, before + // the transferred sub-agent runs, and a later turn resumes at the sub-agent, not the finished + // root. + @Test + @SuppressWarnings("deprecation") // Resumability flag is intentionally partial. + public void runAsync_resumable_transferToSubAgent_closesRootThenResumesSubAgent() { + Content transferCall = + Content.fromParts( + Part.fromFunctionCall( + "transfer_to_agent", ImmutableMap.of("agent_name", "sub_agent_1"))); + TestLlm testLlm = + createTestLlm( + createLlmResponse(transferCall), + createTextLlmResponse("response1"), + createTextLlmResponse("response2")); + LlmAgent subAgent1 = createTestAgentBuilder(testLlm).name("sub_agent_1").build(); + LlmAgent rootAgent = + createTestAgentBuilder(testLlm) + .name("root_agent") + .subAgents(ImmutableList.of(subAgent1)) + .build(); + Runner runner = + Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(rootAgent) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build()) + .build(); + Session session = runner.sessionService().createSession("test", "user").blockingGet(); + + List turn1 = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("hi"))) + .toList() + .blockingGet(); + + // The root closes right after the transfer, before the sub-agent runs. + assertThat(simplifyResumableEvents(turn1)) + .containsExactly( + "root_agent: FunctionCall(name=transfer_to_agent, args={agent_name=sub_agent_1})", + "root_agent: FunctionResponse(name=transfer_to_agent, response={})", + "root_agent: end_of_agent", + "sub_agent_1: response1", + "sub_agent_1: end_of_agent") + .inOrder(); + + List turn2 = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("again"))) + .toList() + .blockingGet(); + + // A new turn resumes at the transferred sub-agent, not the finished root. + assertThat(simplifyEvents(turn2)).contains("sub_agent_1: response2"); + } + + // A sub-agent a transfer routed to can itself pause on a long-running call and be resumed: the + // root closes on transfer, the sub-agent pauses on its long-running call, and a resume carrying + // the matching function response continues that same sub-agent invocation to completion. + @Test + @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). + public void runAsync_resumable_transferredSubAgentPausesOnLongRunningCall_resumesSubAgent() { + Content transferCall = + Content.fromParts( + Part.fromFunctionCall( + "transfer_to_agent", ImmutableMap.of("agent_name", "sub_agent_1"))); + TestLlm testLlm = + createTestLlm( + createLlmResponse(transferCall), + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hello")), + createTextLlmResponse("resumed answer")); + LlmAgent subAgent1 = + createTestAgentBuilder(testLlm) + .name("sub_agent_1") + .tools( + FunctionTool.create( + Tools.class, + "pendingTool", + /* requireConfirmation= */ false, + /* isLongRunning= */ true)) + .build(); + LlmAgent rootAgent = + createTestAgentBuilder(testLlm) + .name("root_agent") + .subAgents(ImmutableList.of(subAgent1)) + .build(); + Runner runner = + Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(rootAgent) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build()) + .build(); + Session session = runner.sessionService().createSession("test", "user").blockingGet(); + + List pausedTurn = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("hi"))) + .toList() + .blockingGet(); + String pausedInvocationId = pausedTurn.get(0).invocationId(); + + // Root closes on transfer, then the sub-agent runs and pauses on its long-running call. + ImmutableList pausedEvents = simplifyResumableEvents(pausedTurn); + assertThat(pausedEvents) + .containsAtLeast( + "root_agent: end_of_agent", + "sub_agent_1: FunctionCall(name=pendingTool, args={message=hello})") + .inOrder(); + // The sub-agent neither finished nor produced its answer while paused. + assertThat(pausedEvents).doesNotContain("sub_agent_1: end_of_agent"); + assertThat(simplifyEvents(pausedTurn)).doesNotContain("sub_agent_1: resumed answer"); + + List resumed = + runner + .runAsync( + "user", + session.id(), + /* invocationId= */ null, + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("lro_call_id") + .name("pendingTool") + .response(ImmutableMap.of("message", "hello"))) + .build()), + RunConfig.builder().build(), + /* stateDelta= */ null) + .toList() + .blockingGet(); + + // The resume continues the transferred sub-agent (same invocation) to completion. + assertThat(resumed).isNotEmpty(); + assertThat(resumed.stream().allMatch(event -> event.invocationId().equals(pausedInvocationId))) + .isTrue(); + assertThat(simplifyEvents(resumed)).contains("sub_agent_1: resumed answer"); + assertThat(resumed.stream().anyMatch(event -> event.actions().endOfAgent())).isTrue(); + } + + // Regression: with the plain-text auto-resume shim, a transferred invocation must not look + // unfinished forever. After a transfer the root closes and later turns run the sub-agent, so a + // finished-check keyed on the root wedged every plain-text turn from turn 3 on. + @Test + @SuppressWarnings("deprecation") // exercises the deprecated plainTextContinuationAutoResume shim + public void runAsync_resumableTransferWithPlainTextAutoResume_laterTurnsRunSubAgent() { + Content transferCall = + Content.fromParts( + Part.fromFunctionCall( + "transfer_to_agent", ImmutableMap.of("agent_name", "sub_agent_1"))); + TestLlm testLlm = + createTestLlm( + createLlmResponse(transferCall), + createTextLlmResponse("r1"), + createTextLlmResponse("r2"), + createTextLlmResponse("r3"), + createTextLlmResponse("r4"), + createTextLlmResponse("r5")); + LlmAgent subAgent1 = createTestAgentBuilder(testLlm).name("sub_agent_1").build(); + LlmAgent rootAgent = + createTestAgentBuilder(testLlm) + .name("root_agent") + .subAgents(ImmutableList.of(subAgent1)) + .build(); + Runner runner = + Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(rootAgent) + .resumabilityConfig( + ResumabilityConfig.builder() + .resumable(true) + .plainTextContinuationAutoResume(true) + .build()) + .build()) + .build(); + Session session = runner.sessionService().createSession("test", "user").blockingGet(); + + // Turn 1 transfers to the sub-agent; turns 2-5 (plain text) must each be answered by the + // sub-agent rather than wedging to an empty stream. + var unused = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("m1"))) + .toList() + .blockingGet(); + for (String expected : new String[] {"r2", "r3", "r4", "r5"}) { + List turn = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("m"))) + .toList() + .blockingGet(); + assertThat(simplifyEvents(turn)).contains("sub_agent_1: " + expected); + } + } + + // Rollout guard: a completed checkpoint-less session (created before checkpoints existed) has no + // end-of-agent signal, so the plain-text auto-resume shim must not re-attach to it; it starts a + // new invocation, keeping the per-turn invocation scoping downstream callbacks rely on. + @Test + @SuppressWarnings("deprecation") // Exercises the deprecated plainTextContinuationAutoResume shim. + public void runAsync_plainTextAutoResume_checkpointlessSession_startsNewInvocation() { + LlmAgent agent = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("second answer"))) + .name("agent") + .build(); + Runner runner = + Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(agent) + .resumabilityConfig( + ResumabilityConfig.builder() + .resumable(true) + .plainTextContinuationAutoResume(true) + .build()) + .build()) + .build(); + Session session = runner.sessionService().createSession("test", "user").blockingGet(); + + // Seed a completed prior turn with NO resumability checkpoints (no endOfAgent / agentState), as + // a pre-checkpoint session looks on the wire. + String priorInvocationId = "pre_checkpoint_invocation"; + Event unusedUserEvent = + runner + .sessionService() + .appendEvent( + session, + Event.builder() + .id(Event.generateEventId()) + .invocationId(priorInvocationId) + .author("user") + .content(Content.fromParts(Part.fromText("first turn"))) + .build()) + .blockingGet(); + Event unusedModelEvent = + runner + .sessionService() + .appendEvent( + session, + Event.builder() + .id(Event.generateEventId()) + .invocationId(priorInvocationId) + .author("agent") + .content(Content.fromParts(Part.fromText("first answer"))) + .build()) + .blockingGet(); + + List secondTurn = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("second turn"))) + .toList() + .blockingGet(); + + // The shim must not re-attach to the checkpoint-less prior invocation: the agent runs under a + // fresh invocation id. + assertThat(simplifyEvents(secondTurn)).contains("agent: second answer"); + assertThat( + secondTurn.stream() + .map(Event::invocationId) + .filter(priorInvocationId::equals) + .collect(toImmutableList())) + .isEmpty(); + } + + // A resumable invocation paused on two long-running calls stays paused until both are answered: + // answering one resumes without re-invoking the model, and answering the second lets the model + // summarize. + @Test + @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). + public void runAsync_withTwoLongRunningCalls_pausesUntilBothAnswered() { + TestLlm testLlm = + createTestLlm( + createLlmResponse( + Content.builder() + .role("model") + .parts( + Part.builder() + .functionCall( + FunctionCall.builder() + .id("call_a") + .name("pendingTool") + .args(ImmutableMap.of("message", "a"))) + .build(), + Part.builder() + .functionCall( + FunctionCall.builder() + .id("call_b") + .name("pendingTool") + .args(ImmutableMap.of("message", "b"))) + .build()) + .build()), + createTextLlmResponse("both approved")); + LlmAgent agent = + createTestAgentBuilder(testLlm) + .name("root_agent") + .tools( + FunctionTool.create( + Tools.class, + "pendingTool", + /* requireConfirmation= */ false, + /* isLongRunning= */ true)) + .build(); + Runner runner = + Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(agent) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build()) + .build(); + Session session = runner.sessionService().createSession("test", "user").blockingGet(); + + Object unusedFirst = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("start"))) + .toList() + .blockingGet(); + // Turn 1: both long-running calls are issued and the invocation pauses; the model is called + // once. + assertThat(testLlm.getRequests()).hasSize(1); + + List afterFirstAnswer = + runner + .runAsync( + "user", + session.id(), + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("call_a") + .name("pendingTool") + .response(ImmutableMap.of("message", "a"))) + .build())) + .toList() + .blockingGet(); + // One call answered is not enough: nothing runs and the model is not re-invoked. + assertThat(afterFirstAnswer).isEmpty(); + assertThat(testLlm.getRequests()).hasSize(1); + + List afterSecondAnswer = + runner + .runAsync( + "user", + session.id(), + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("call_b") + .name("pendingTool") + .response(ImmutableMap.of("message", "b"))) + .build())) + .toList() + .blockingGet(); + // Both answered: the model is re-invoked and summarizes. + assertThat(testLlm.getRequests()).hasSize(2); + assertThat(simplifyEvents(afterSecondAnswer)).contains("root_agent: both approved"); + } + + // A value-returning long-running tool is not a pending request: it resolves the call in the same + // turn, so even with resumability on the flow continues and the model summarizes the result (two + // model calls) rather than pausing. Only a no-result long-running tool pauses. + @Test + @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). + public void runAsync_withLongRunningCall_resumable_valueReturn_continuesAndSummarizes() { + TestLlm testLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "echoTool", ImmutableMap.of("message", "hello")), + createTextLlmResponse("summarized echo")); + LlmAgent agent = + createTestAgentBuilder(testLlm) + .name("agent") + .tools( + FunctionTool.create( + Tools.class, + "echoTool", + /* requireConfirmation= */ false, + /* isLongRunning= */ true)) + .build(); + Runner runner = + Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(agent) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build()) + .build(); + Session session = runner.sessionService().createSession("test", "user").blockingGet(); + + List events = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) + .toList() + .blockingGet(); + + // The value-returning call was summarized in the same turn: the model was re-invoked (two + // calls) and the summary surfaced, with no pause. + assertThat(testLlm.getRequests()).hasSize(2); + assertThat(simplifyEvents(events)).contains("agent: summarized echo"); + } + + // Pin: a value-returning long-running tool is summarized but emits no end-of-agent checkpoint, so + // the invocation stays resumable. + @Test + @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). + public void runAsync_resumable_valueReturnLongRunning_emitsNoEndOfAgentCheckpoint() { + TestLlm testLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "echoTool", ImmutableMap.of("message", "hello")), + createTextLlmResponse("summarized echo")); + LlmAgent agent = + createTestAgentBuilder(testLlm) + .name("agent") + .tools( + FunctionTool.create( + Tools.class, + "echoTool", + /* requireConfirmation= */ false, + /* isLongRunning= */ true)) + .build(); + Runner runner = + Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(agent) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build()) + .build(); + Session session = runner.sessionService().createSession("test", "user").blockingGet(); + + List events = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) + .toList() + .blockingGet(); + + assertThat(simplifyEvents(events)).contains("agent: summarized echo"); + assertThat(events.stream().anyMatch(event -> event.actions().endOfAgent())).isFalse(); + } + + // On resume the runner runs the same plugin bracket as the new-invocation path (on-user-message, + // before-run, after-run, on-event), not only on-event: each fires once on the initial turn and + // once more on the resume turn. + @Test + @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). + public void runAsync_resume_runsFullPluginBracket() { + BasePlugin resumePlugin = mockPlugin("resume"); + TestLlm testLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hello")), + createTextLlmResponse("resumed and summarized")); + LlmAgent agent = + createTestAgentBuilder(testLlm) + .name("agent") + .tools( + FunctionTool.create( + Tools.class, + "pendingTool", + /* requireConfirmation= */ false, + /* isLongRunning= */ true)) + .build(); + Runner runner = + Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(agent) + .plugins(ImmutableList.of(resumePlugin)) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build()) + .build(); + Session session = runner.sessionService().createSession("test", "user").blockingGet(); + + Object unusedFirstTurn = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("start"))) + .toList() + .blockingGet(); + + List resumed = + runner + .runAsync( + "user", + session.id(), + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("lro_call_id") + .name("pendingTool") + .response(ImmutableMap.of("message", "hello"))) + .build())) + .toList() + .blockingGet(); + + assertThat(simplifyEvents(resumed)).contains("agent: resumed and summarized"); + // The full bracket ran on both the initial turn and the resume turn (before the fix, the resume + // turn ran only onEventCallback, so these would each be invoked once). + verify(resumePlugin, times(2)).onUserMessageCallback(any(), any()); + verify(resumePlugin, times(2)).beforeRunCallback(any()); + verify(resumePlugin, times(2)).afterRunCallback(any()); + verify(resumePlugin, atLeastOnce()).onEventCallback(any(), any()); + } + + // A resumable LlmAgent that has a checkpoint, resumed with a plain-text message (not a function + // response), continues the agent instead of throwing "No matching function call". The + // still-unanswered long-running call keeps the invocation paused, so nothing is emitted. + @Test + @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). + public void resume_withPlainTextMessage_continuesWithoutThrowing() { + TestLlm testLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hello")), + createTextLlmResponse("should not be reached")); + LlmAgent agent = + createTestAgentBuilder(testLlm) + .name("agent") + .tools( + FunctionTool.create( + Tools.class, + "pendingTool", + /* requireConfirmation= */ false, + /* isLongRunning= */ true)) + .build(); + Runner runner = + Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(agent) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build()) + .build(); + Session session = runner.sessionService().createSession("test", "user").blockingGet(); + + List firstTurn = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("start"))) + .toList() + .blockingGet(); + String invocationId = firstTurn.get(0).invocationId(); + + // Resume the paused invocation with a plain-text message and an explicit invocation id. Before + // the fix this threw IllegalArgumentException; now it returns without throwing. + List resumed = + runner + .runAsync( + "user", + session.id(), + invocationId, + Content.fromParts(Part.fromText("please continue")), + RunConfig.builder().build(), + /* stateDelta= */ null) + .toList() + .blockingGet(); + + assertThat(resumed).isEmpty(); + } + + // Pin: a leaf paused under a ParallelAgent resumes alone; the enclosing SequentialAgent does + // not advance (ParallelAgent isn't resume-aware). + @Test + @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). + public void runAsync_resume_pausedUnderParallelAgent_doesNotAdvanceEnclosingSequential() { + LlmAgent leaf = + createTestAgentBuilder( + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hi")), + createTextLlmResponse("leaf resumed"))) + .name("leaf_agent") + .tools( + FunctionTool.create( + Tools.class, + "pendingTool", + /* requireConfirmation= */ false, + /* isLongRunning= */ true)) + .build(); + ParallelAgent parallel = + ParallelAgent.builder().name("parallel_agent").subAgents(ImmutableList.of(leaf)).build(); + LlmAgent next = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("next ran"))) + .name("next_agent") + .build(); + SequentialAgent root = + SequentialAgent.builder() + .name("seq_agent") + .subAgents(ImmutableList.of(parallel, next)) + .build(); + Runner runner = + Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(root) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build()) + .build(); + Session session = runner.sessionService().createSession("test", "user").blockingGet(); + + Object unusedFirstTurn = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("start"))) + .toList() + .blockingGet(); + + List resumed = + runner + .runAsync( + "user", + session.id(), + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("lro_call_id") + .name("pendingTool") + .response(ImmutableMap.of("message", "hi"))) + .build())) + .toList() + .blockingGet(); + + // The leaf itself resumes and emits its post-tool response. + assertThat(simplifyEvents(resumed)).contains("leaf_agent: leaf resumed"); + // But the enclosing SequentialAgent does not advance: next_agent never runs. + assertThat(simplifyEvents(resumed)).doesNotContain("next_agent: next ran"); + } + + // A leaf paused on two long-running calls under a ParallelAgent, answered once, stays paused + // (parent-branch seeding keeps the other call visible). + @Test + @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). + public void runAsync_resume_pausedUnderParallelAgent_partiallyAnswered_staysPaused() { + Content twoLongRunningCalls = + Content.builder() + .role("model") + .parts( + Part.builder() + .functionCall( + FunctionCall.builder() + .id("c1") + .name("pendingTool") + .args(ImmutableMap.of("message", "a")) + .build()) + .build(), + Part.builder() + .functionCall( + FunctionCall.builder() + .id("c2") + .name("pendingTool") + .args(ImmutableMap.of("message", "b")) + .build()) + .build()) + .build(); + LlmAgent leaf = + createTestAgentBuilder( + createTestLlm( + createLlmResponse(twoLongRunningCalls), createTextLlmResponse("leaf summary"))) + .name("leaf_agent") + .tools( + FunctionTool.create( + Tools.class, + "pendingTool", + /* requireConfirmation= */ false, + /* isLongRunning= */ true)) + .build(); + ParallelAgent parallel = + ParallelAgent.builder().name("parallel_agent").subAgents(ImmutableList.of(leaf)).build(); + SequentialAgent root = + SequentialAgent.builder().name("seq_agent").subAgents(ImmutableList.of(parallel)).build(); + Runner runner = + Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(root) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build()) + .build(); + Session session = runner.sessionService().createSession("test", "user").blockingGet(); + + Object unusedFirstTurn = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("start"))) + .toList() + .blockingGet(); + + // Answer only c1; c2 remains unanswered. + List resumed = + runner + .runAsync( + "user", + session.id(), + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("c1") + .name("pendingTool") + .response(ImmutableMap.of("message", "a"))) + .build())) + .toList() + .blockingGet(); + + // c2 is still unanswered, so the model is not re-invoked: no "leaf summary", no new events. + assertThat(simplifyEvents(resumed)).doesNotContain("leaf_agent: leaf summary"); + assertThat(resumed).isEmpty(); + } + + // Opt-in shim: with plainTextContinuationAutoResume(true), a plain-text continuation resumes the + // paused invocation instead of starting a new one. + @Test + @SuppressWarnings( + "deprecation") // Resumability + the auto-resume shim are intentionally deprecated. + public void runAsync_plainTextContinuation_autoResumeFlagOn_resumesPausedInvocation() { + TestLlm testLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "draft")), + createTextLlmResponse("should not re-plan")); + LlmAgent agent = + createTestAgentBuilder(testLlm) + .name("agent") + .tools( + FunctionTool.create( + Tools.class, + "pendingTool", + /* requireConfirmation= */ false, + /* isLongRunning= */ true)) + .build(); + Runner runner = + Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(agent) + .resumabilityConfig( + ResumabilityConfig.builder() + .resumable(true) + .plainTextContinuationAutoResume(true) + .build()) + .build()) + .build(); + Session session = runner.sessionService().createSession("test", "user").blockingGet(); + + List turn1 = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("draft the note"))) + .toList() + .blockingGet(); + String invocationId = turn1.get(0).invocationId(); + assertThat(testLlm.getRequests()).hasSize(1); // paused after a single model call + + List turn2 = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("Proceed"))) + .toList() + .blockingGet(); + + assertThat(testLlm.getRequests()).hasSize(1); // resumed and held: no re-plan + assertThat(turn2).isEmpty(); + Session reloaded = + runner + .sessionService() + .getSession("test", "user", session.id(), Optional.empty()) + .blockingGet(); + Event lastUserEvent = + Streams.findLast( + reloaded.events().stream().filter(event -> Objects.equals(event.author(), "user"))) + .orElse(null); + assertThat(lastUserEvent).isNotNull(); + assertThat(lastUserEvent.invocationId()).isEqualTo(invocationId); + } + + // Guard: even with the auto-resume shim on, a plain-text message after a finished turn starts a + // NEW invocation (the shim resumes only unfinished invocations, so it never swallows a new turn). + @Test + @SuppressWarnings( + "deprecation") // Resumability + the auto-resume shim are intentionally deprecated. + public void runAsync_plainText_autoResumeFlagOn_afterCompletedTurn_startsNewInvocation() { + TestLlm testLlm = + createTestLlm( + createTextLlmResponse("first answer"), createTextLlmResponse("second answer")); + LlmAgent agent = createTestAgentBuilder(testLlm).name("agent").build(); + Runner runner = + Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(agent) + .resumabilityConfig( + ResumabilityConfig.builder() + .resumable(true) + .plainTextContinuationAutoResume(true) + .build()) + .build()) + .build(); + Session session = runner.sessionService().createSession("test", "user").blockingGet(); + + List turn1 = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("hi"))) + .toList() + .blockingGet(); + String invocationId = turn1.get(0).invocationId(); + assertThat(simplifyEvents(turn1)).contains("agent: first answer"); + + List turn2 = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("again"))) + .toList() + .blockingGet(); + + assertThat(simplifyEvents(turn2)).contains("agent: second answer"); + assertThat(turn2.get(0).invocationId()).isNotEqualTo(invocationId); + } + + // Default (shim off): a plain-text continuation after a pause starts a NEW invocation, not a + // resume; resuming is explicit (a function response or runAsync with an invocation id). + @Test + @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). + public void runAsync_plainTextContinuation_autoResumeFlagOff_startsNewInvocation() { + TestLlm testLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "draft")), + createTextLlmResponse("re-planned")); + LlmAgent agent = + createTestAgentBuilder(testLlm) + .name("agent") + .tools( + FunctionTool.create( + Tools.class, + "pendingTool", + /* requireConfirmation= */ false, + /* isLongRunning= */ true)) + .build(); + Runner runner = + Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(agent) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build()) + .build(); + Session session = runner.sessionService().createSession("test", "user").blockingGet(); + + List turn1 = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("draft the note"))) + .toList() + .blockingGet(); + String invocationId = turn1.get(0).invocationId(); + + List turn2 = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("Proceed"))) + .toList() + .blockingGet(); + + assertThat(testLlm.getRequests()).hasSize(2); // new invocation re-invoked the model + assertThat(turn2.get(0).invocationId()).isNotEqualTo(invocationId); + } + + // Nested topology: the plain-text auto-resume shim also resumes a paused invocation when the + // paused long-running call sits inside an LlmAgent nested in a SequentialAgent (as in the + // HITL-in-SequentialAgent bug), not just a flat agent, instead of starting a new invocation. + @Test + @SuppressWarnings( + "deprecation") // Resumability + the auto-resume shim are intentionally deprecated. + public void + runAsync_plainTextContinuation_inSequentialAgent_autoResumeFlagOn_resumesPausedInvocation() { + TestLlm testLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "draft")), + createTextLlmResponse("should not re-plan")); + LlmAgent childAgent = + createTestAgentBuilder(testLlm) + .name("child_agent") + .tools( + FunctionTool.create( + Tools.class, + "pendingTool", + /* requireConfirmation= */ false, + /* isLongRunning= */ true)) + .build(); + SequentialAgent workflowAgent = + SequentialAgent.builder() + .name("workflow_agent") + .subAgents(ImmutableList.of(childAgent)) + .build(); + Runner runner = + Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(workflowAgent) + .resumabilityConfig( + ResumabilityConfig.builder() + .resumable(true) + .plainTextContinuationAutoResume(true) + .build()) + .build()) + .build(); + Session session = runner.sessionService().createSession("test", "user").blockingGet(); + + List turn1 = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("draft the note"))) + .toList() + .blockingGet(); + String invocationId = turn1.get(0).invocationId(); + assertThat(testLlm.getRequests()).hasSize(1); // paused inside the workflow after one model call + + List turn2 = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("Proceed"))) + .toList() + .blockingGet(); + + // The shim resumed the paused invocation instead of re-planning a new one: no extra model call + // (so the second scripted response is never reached), and the appended user turn carries the + // paused invocation id. + assertThat(testLlm.getRequests()).hasSize(1); + assertThat(simplifyEvents(turn2)).doesNotContain("child_agent: should not re-plan"); + Session reloaded = + runner + .sessionService() + .getSession("test", "user", session.id(), Optional.empty()) + .blockingGet(); + Event lastUserEvent = + Streams.findLast( + reloaded.events().stream().filter(event -> Objects.equals(event.author(), "user"))) + .orElse(null); + assertThat(lastUserEvent).isNotNull(); + assertThat(lastUserEvent.invocationId()).isEqualTo(invocationId); + } + + // Gating: with resumability OFF (default) the flow does NOT pause on a long-running call; it + // keeps + // calling the model as before. Pairs with the resumable test above. + @Test + public void runAsync_withLongRunningCall_resumabilityDisabled_doesNotPause() { + TestLlm testLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "echoTool", ImmutableMap.of("message", "hello")), + createTextLlmResponse("after pending call")); + LlmAgent agent = + createTestAgentBuilder(testLlm) + .name("agent") + .tools( + FunctionTool.create( + Tools.class, + "echoTool", + /* requireConfirmation= */ false, + /* isLongRunning= */ true)) + .build(); + Runner runner = + Runner.builder().app(App.builder().name("test").rootAgent(agent).build()).build(); + Session session = runner.sessionService().createSession("test", "user").blockingGet(); + + List events = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) + .toList() + .blockingGet(); + + // No pause: the flow made a second model call and surfaced its response. + assertThat(testLlm.getRequests()).hasSize(2); + assertThat(simplifyEvents(events)).contains("agent: after pending call"); + } + + // A long-running tool awaiting an external result (real HITL, e.g. human input) returns nothing + // yet. The invocation must end after the single model call rather than re-invoking the model with + // a placeholder response and looping until the call limit. Matches Python ADK v1: the function + // response is skipped and the long-running call event is treated as final. + @Test + public void runAsync_withLongRunningCall_noImmediateResult_endsAfterSingleModelCall() { + TestLlm testLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hello")), + // Extra response the flow must NOT consume; reaching it means it looped. + createTextLlmResponse("should not be reached")); + LlmAgent agent = + createTestAgentBuilder(testLlm) + .name("agent") + .tools( + FunctionTool.create( + Tools.class, + "pendingTool", + /* requireConfirmation= */ false, + /* isLongRunning= */ true)) + .build(); + Runner runner = + Runner.builder().app(App.builder().name("test").rootAgent(agent).build()).build(); + Session session = runner.sessionService().createSession("test", "user").blockingGet(); + + List events = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) + .toList() + .blockingGet(); + + // Ended after the single long-running call: no function response, no second model call. + assertThat(testLlm.getRequests()).hasSize(1); + assertThat(simplifyEvents(events)).doesNotContain("agent: should not be reached"); + } + + // A resumable LlmAgent that completes normally emits a trailing end-of-agent checkpoint so a + // later run can tell the invocation finished. + @Test + @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). + public void runAsync_resumable_completedLlmAgent_emitsEndOfAgentCheckpoint() { + LlmAgent agent = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("all done"))) + .name("agent") + .build(); + Runner runner = + Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(agent) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build()) + .build(); + Session session = runner.sessionService().createSession("test", "user").blockingGet(); + + List events = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) + .toList() + .blockingGet(); + + Event last = Iterables.getLast(events); + assertThat(last.author()).isEqualTo("agent"); + assertThat(last.actions().endOfAgent()).isTrue(); + } + + // Gating: with resumability OFF (default) a completed LlmAgent emits no end-of-agent checkpoint, + // keeping the event stream identical to before. Pairs with the resumable test above. + @Test + public void runAsync_resumabilityDisabled_completedLlmAgent_emitsNoEndOfAgent() { + LlmAgent agent = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("all done"))) + .name("agent") + .build(); + Runner runner = + Runner.builder().app(App.builder().name("test").rootAgent(agent).build()).build(); + Session session = runner.sessionService().createSession("test", "user").blockingGet(); + + List events = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) + .toList() + .blockingGet(); + + assertThat(events.stream().anyMatch(event -> event.actions().endOfAgent())).isFalse(); + } + + // Resuming a completed invocation is a no-op: the active agent already ended, so nothing + // re-runs. + @Test + @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). + public void resume_completedInvocation_isNoOp() { + LlmAgent agent = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("all done"))) + .name("agent") + .build(); + Runner runner = + Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(agent) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build()) + .build(); + Session session = runner.sessionService().createSession("test", "user").blockingGet(); + + List firstTurn = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) + .toList() + .blockingGet(); + String invocationId = firstTurn.get(0).invocationId(); + + List resumed = + runner + .runAsync( + "user", + session.id(), + invocationId, + /* newMessage= */ null, + RunConfig.builder().build(), + /* stateDelta= */ null) + .toList() + .blockingGet(); + + assertThat(resumed).isEmpty(); + } + + // Resuming a paused long-running call WITHOUT an answer stays paused: the model must not be + // re-invoked while the call is unanswered. + @Test + @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). + public void resume_pausedCallWithoutAnswer_staysPaused() { + TestLlm testLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hello")), + createTextLlmResponse("resumed answer")); + LlmAgent agent = + createTestAgentBuilder(testLlm) + .name("agent") + .tools( + FunctionTool.create( + Tools.class, + "pendingTool", + /* requireConfirmation= */ false, + /* isLongRunning= */ true)) + .build(); + Runner runner = + Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(agent) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build()) + .build(); + Session session = runner.sessionService().createSession("test", "user").blockingGet(); + + List pausedTurn = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) + .toList() + .blockingGet(); + String pausedInvocationId = pausedTurn.get(0).invocationId(); + assertThat(simplifyEvents(pausedTurn)).doesNotContain("agent: resumed answer"); + + // Resume WITHOUT providing the function response: the paused call is still unanswered. + List resumed = + runner + .runAsync( + "user", + session.id(), + pausedInvocationId, + /* newMessage= */ null, + RunConfig.builder().build(), + /* stateDelta= */ null) + .toList() + .blockingGet(); + + // Still unanswered, so the model is not re-invoked and no new content is produced. + assertThat(simplifyEvents(resumed)).doesNotContain("agent: resumed answer"); + } + + // Resuming with a function response resumes the SAME invocation that issued the matching call + // (rather than minting a new invocation id) and runs it to completion. + @Test + @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). + public void resume_withFunctionResponse_resumesSameInvocation() { + TestLlm testLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hello")), + createTextLlmResponse("resumed answer")); + LlmAgent agent = + createTestAgentBuilder(testLlm) + .name("agent") + .tools( + FunctionTool.create( + Tools.class, + "pendingTool", + /* requireConfirmation= */ false, + /* isLongRunning= */ true)) + .build(); + Runner runner = + Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(agent) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build()) + .build(); + Session session = runner.sessionService().createSession("test", "user").blockingGet(); + + List pausedTurn = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) + .toList() + .blockingGet(); + String pausedInvocationId = pausedTurn.get(0).invocationId(); + assertThat(simplifyEvents(pausedTurn)).doesNotContain("agent: resumed answer"); + + List resumed = + runner + .runAsync( + "user", + session.id(), + /* invocationId= */ null, + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("lro_call_id") + .name("pendingTool") + .response(ImmutableMap.of("message", "hello"))) + .build()), + RunConfig.builder().build(), + /* stateDelta= */ null) + .toList() + .blockingGet(); + + // The resumed events belong to the original (paused) invocation, not a fresh one. + assertThat(resumed).isNotEmpty(); + assertThat(resumed.stream().allMatch(event -> event.invocationId().equals(pausedInvocationId))) + .isTrue(); + assertThat(simplifyEvents(resumed)).contains("agent: resumed answer"); + assertThat(resumed.stream().anyMatch(event -> event.actions().endOfAgent())).isTrue(); + } + + // The resume overload merges a non-null stateDelta into the session, like the new-invocation + // path. + @Test + @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). + public void resume_withStateDelta_mergesStateIntoSession() { + TestLlm testLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hello")), + createTextLlmResponse("resumed answer")); + LlmAgent agent = + createTestAgentBuilder(testLlm) + .name("agent") + .tools( + FunctionTool.create( + Tools.class, + "pendingTool", + /* requireConfirmation= */ false, + /* isLongRunning= */ true)) + .build(); + Runner runner = + Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(agent) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build()) + .build(); + Session session = runner.sessionService().createSession("test", "user").blockingGet(); + + Object unusedPausedTurn = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) + .toList() + .blockingGet(); + + ImmutableMap stateDelta = ImmutableMap.of("key1", "value1", "key2", 42); + List resumed = + runner + .runAsync( + "user", + session.id(), + /* invocationId= */ null, + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("lro_call_id") + .name("pendingTool") + .response(ImmutableMap.of("message", "hello"))) + .build()), + RunConfig.builder().build(), + stateDelta) + .toList() + .blockingGet(); + + assertThat(simplifyEvents(resumed)).contains("agent: resumed answer"); + Session finalSession = + runner + .sessionService() + .getSession("test", "user", session.id(), Optional.empty()) + .blockingGet(); + assertThat(finalSession.state()).containsAtLeastEntriesIn(stateDelta); + // The delta is also stamped on the appended (function-response) event, for history rehydration. + Event lastUserEvent = + Streams.findLast( + finalSession.events().stream() + .filter(event -> Objects.equals(event.author(), "user"))) + .orElseThrow(); + assertThat(lastUserEvent.actions().stateDelta()).containsAtLeastEntriesIn(stateDelta); + } + + // With the deprecated plain-text auto-resume flag, a plain-text continuation resumes the last + // unfinished invocation and a non-null stateDelta is still merged. + @Test + @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). + public void resume_plainTextAutoResume_withStateDelta_mergesStateIntoSession() { + TestLlm testLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hello")), + createTextLlmResponse("should not be reached")); + LlmAgent agent = + createTestAgentBuilder(testLlm) + .name("agent") + .tools( + FunctionTool.create( + Tools.class, + "pendingTool", + /* requireConfirmation= */ false, + /* isLongRunning= */ true)) + .build(); + Runner runner = + Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(agent) + .resumabilityConfig( + ResumabilityConfig.builder() + .resumable(true) + .plainTextContinuationAutoResume(true) + .build()) + .build()) + .build(); + Session session = runner.sessionService().createSession("test", "user").blockingGet(); + + List pausedTurn = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("start"))) + .toList() + .blockingGet(); + String pausedInvocationId = pausedTurn.get(0).invocationId(); + assertThat(testLlm.getRequests()).hasSize(1); // paused after a single model call + + // Plain-text "Proceed" via the non-resume overload; the flag resumes the paused invocation. + ImmutableMap stateDelta = ImmutableMap.of("key1", "value1", "key2", 42); + List resumed = + runner + .runAsync( + "user", + session.id(), + Content.fromParts(Part.fromText("Proceed")), + RunConfig.builder().build(), + stateDelta) + .toList() + .blockingGet(); + + // Resumed and held on the unanswered long-running call: no re-plan, no new invocation. + assertThat(testLlm.getRequests()).hasSize(1); + assertThat(resumed).isEmpty(); + Session finalSession = + runner + .sessionService() + .getSession("test", "user", session.id(), Optional.empty()) + .blockingGet(); + Event continuation = + Streams.findLast( + finalSession.events().stream() + .filter(event -> Objects.equals(event.author(), "user"))) + .orElseThrow(); + assertThat(continuation.invocationId()).isEqualTo(pausedInvocationId); + assertThat(continuation.actions().stateDelta()).containsAtLeastEntriesIn(stateDelta); + assertThat(finalSession.state()).containsAtLeastEntriesIn(stateDelta); + } + + // ResumeInvocationTest parity: resume an OLDER paused invocation (not the latest) via its + // long-running function response; the resumed run belongs to that older invocation. + @Test + @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). + public void resume_resumesAnyInvocation_notJustTheLatest() { + TestLlm testLlm = + createTestLlm( + createFunctionCallLlmResponse( + "call-1", "pendingTool", ImmutableMap.of("message", "hi")), + createTextLlmResponse("llm response in invocation 2"), + createFunctionCallLlmResponse( + "call-3", "pendingTool", ImmutableMap.of("message", "hi")), + createTextLlmResponse("llm response after resuming invocation 1")); + LlmAgent agent = + createTestAgentBuilder(testLlm) + .name("agent") .tools( FunctionTool.create( Tools.class, - "echoTool", + "pendingTool", /* requireConfirmation= */ false, /* isLongRunning= */ true)) .build(); - LlmAgent agentC = - createTestAgentBuilder(createTestLlm(createTextLlmResponse("agent C done"))) - .name("c_agent") - .build(); - SequentialAgent workflowAgent = - SequentialAgent.builder() - .name("workflow_agent") - .subAgents(ImmutableList.of(agentA, agentB, agentC)) - .build(); Runner runner = Runner.builder() .app( App.builder() .name("test") - .rootAgent(workflowAgent) + .rootAgent(agent) .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) .build()) .build(); Session session = runner.sessionService().createSession("test", "user").blockingGet(); - List eventsBeforeResume = + // Invocation 1 pauses on the long-running call. + List inv1 = runner - .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) + .runAsync("user", session.id(), Content.fromParts(Part.fromText("q1"))) + .toList() + .blockingGet(); + String inv1Id = inv1.get(0).invocationId(); + // Invocation 2 finishes; invocation 3 pauses again. + Object unusedInv2 = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("q2"))) + .toList() + .blockingGet(); + Object unusedInv3 = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("q3"))) .toList() .blockingGet(); - // Turn 1: A runs, B issues the long-running call and pauses; C must not run yet. B must not - // make - // a further model call after the pending call. - assertThat(simplifyEvents(eventsBeforeResume)).contains("a_agent: agent A done"); - assertThat(simplifyEvents(eventsBeforeResume)).doesNotContain("b_agent: agent B resumed"); - assertThat(simplifyEvents(eventsBeforeResume)).doesNotContain("c_agent: agent C done"); - - List eventsAfterResume = + // Resume invocation 1 (the oldest, not the latest) via its function response. + List resumed = runner .runAsync( "user", session.id(), + inv1Id, Content.fromParts( Part.builder() .functionResponse( FunctionResponse.builder() - .id("lro_call_id") - .name("echoTool") - .response(ImmutableMap.of("message", "hello"))) - .build())) + .id("call-1") + .name("pendingTool") + .response(ImmutableMap.of("message", "hi"))) + .build()), + RunConfig.builder().build(), + /* stateDelta= */ null) .toList() .blockingGet(); - // Turn 2: B resumes from the long-running response, then C runs. A is not re-run. - assertThat(simplifyEvents(eventsAfterResume)) - .containsExactly("b_agent: agent B resumed", "c_agent: agent C done") - .inOrder(); + assertThat(simplifyEvents(resumed)).contains("agent: llm response after resuming invocation 1"); + assertThat(resumed.stream().allMatch(event -> event.invocationId().equals(inv1Id))).isTrue(); } - // Regression: a pending long-running call must pause the LLM flow after a single model call when - // resumability is on. Before the flow-level pause, the flow kept re-calling the model (re-issuing - // the call), burning tokens. The scripted model would re-issue the call if the flow did not - // pause; - // we assert exactly one model call was made and the later responses were never consumed. + // InMemoryRunnerTest parity: resume by invocationId rehydrates the agent's checkpoint state from + // history so the running agent observes it. @Test @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). - public void runAsync_withLongRunningCall_resumable_pausesAfterSingleModelCall() { - TestLlm testLlm = - createTestLlm( - createFunctionCallLlmResponse( - "lro_call_id", "echoTool", ImmutableMap.of("message", "hello")), - // Extra responses the flow must NOT consume; reaching them means it looped. - createFunctionCallLlmResponse( - "lro_call_id", "echoTool", ImmutableMap.of("message", "hello")), - createTextLlmResponse("should not be reached")); - LlmAgent agent = - createTestAgentBuilder(testLlm) - .name("agent") - .tools( - FunctionTool.create( - Tools.class, - "echoTool", - /* requireConfirmation= */ false, - /* isLongRunning= */ true)) - .build(); + public void resume_restoresAgentStateFromHistory() { + TestBaseAgent agent = + new TestBaseAgent( + "test_agent", + "desc", + () -> Flowable.empty(), + /* subAgents= */ null, + /* beforeAgentCallbacks= */ null, + /* afterAgentCallbacks= */ null); Runner runner = Runner.builder() .app( @@ -2554,88 +4107,308 @@ public void runAsync_withLongRunningCall_resumable_pausesAfterSingleModelCall() .build()) .build(); Session session = runner.sessionService().createSession("test", "user").blockingGet(); + Object unusedUser = + runner + .sessionService() + .appendEvent( + session, + Event.builder() + .id("u1") + .invocationId("test-inv") + .author("user") + .content(Content.fromParts(Part.fromText("hi"))) + .build()) + .blockingGet(); + Object unusedState = + runner + .sessionService() + .appendEvent( + session, + Event.builder() + .id("s1") + .invocationId("test-inv") + .author("test_agent") + .actions( + EventActions.builder() + .agentState(ImmutableMap.of("saved", "state")) + .build()) + .content(Content.fromParts(Part.fromText("previous response"))) + .build()) + .blockingGet(); - List events = + Object unused = runner - .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) + .runAsync( + "user", + session.id(), + "test-inv", + /* newMessage= */ null, + RunConfig.builder().build(), + /* stateDelta= */ null) .toList() .blockingGet(); - // The flow paused after the single long-running call instead of re-calling the model. - assertThat(testLlm.getRequests()).hasSize(1); - assertThat(simplifyEvents(events)).doesNotContain("agent: should not be reached"); + assertThat(agent.getLastInvocationContext().agentStates()) + .containsEntry("test_agent", ImmutableMap.of("saved", "state")); } - // Gating: with resumability OFF (default) the flow does NOT pause on a long-running call; it - // keeps - // calling the model as before. Pairs with the resumable test above. + // InMemoryRunnerTest parity: resume by invocationId with a new user message appends that content + // under the resumed invocation. @Test - public void runAsync_withLongRunningCall_resumabilityDisabled_doesNotPause() { - TestLlm testLlm = - createTestLlm( - createFunctionCallLlmResponse( - "lro_call_id", "echoTool", ImmutableMap.of("message", "hello")), - createTextLlmResponse("after pending call")); - LlmAgent agent = - createTestAgentBuilder(testLlm) - .name("agent") - .tools( - FunctionTool.create( - Tools.class, - "echoTool", - /* requireConfirmation= */ false, - /* isLongRunning= */ true)) - .build(); + @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). + public void resume_withNewMessage_appendsUserContentUnderResumedInvocation() { + TestBaseAgent agent = + new TestBaseAgent("test_agent", "desc", () -> Flowable.empty(), null, null, null); Runner runner = - Runner.builder().app(App.builder().name("test").rootAgent(agent).build()).build(); + Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(agent) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build()) + .build(); Session session = runner.sessionService().createSession("test", "user").blockingGet(); + Object unusedUser = + runner + .sessionService() + .appendEvent( + session, + Event.builder() + .id("u1") + .invocationId("test-inv") + .author("user") + .content(Content.fromParts(Part.fromText("hi"))) + .build()) + .blockingGet(); - List events = + Object unused = runner - .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) + .runAsync( + "user", + session.id(), + "test-inv", + Content.fromParts(Part.fromText("New message")), + RunConfig.builder().build(), + /* stateDelta= */ null) .toList() .blockingGet(); - // No pause: the flow made a second model call and surfaced its response. - assertThat(testLlm.getRequests()).hasSize(2); - assertThat(simplifyEvents(events)).contains("agent: after pending call"); + Session reloaded = + runner + .sessionService() + .getSession("test", "user", session.id(), Optional.empty()) + .blockingGet(); + assertThat(reloaded.events()).hasSize(2); + assertThat( + Iterables.getLast(reloaded.events()) + .content() + .flatMap(Content::parts) + .get() + .get(0) + .text()) + .hasValue("New message"); } - // A long-running tool awaiting an external result (real HITL, e.g. human input) returns nothing - // yet. The invocation must end after the single model call rather than re-invoking the model with - // a placeholder response and looping until the call limit. Matches Python ADK v1: the function - // response is skipped and the long-running call event is treated as final. + // RunnerTest parity (disabled counterpart): resuming a non-resumable app throws. @Test - public void runAsync_withLongRunningCall_noImmediateResult_endsAfterSingleModelCall() { - TestLlm testLlm = - createTestLlm( - createFunctionCallLlmResponse( - "lro_call_id", "pendingTool", ImmutableMap.of("message", "hello")), - // Extra response the flow must NOT consume; reaching it means it looped. - createTextLlmResponse("should not be reached")); + public void resume_notResumable_throwsException() { LlmAgent agent = - createTestAgentBuilder(testLlm) - .name("agent") - .tools( - FunctionTool.create( - Tools.class, - "pendingTool", - /* requireConfirmation= */ false, - /* isLongRunning= */ true)) - .build(); + createTestAgentBuilder(createTestLlm(createTextLlmResponse("x"))).name("agent").build(); Runner runner = Runner.builder().app(App.builder().name("test").rootAgent(agent).build()).build(); Session session = runner.sessionService().createSession("test", "user").blockingGet(); + String sessionId = session.id(); - List events = + RunConfig runConfig = RunConfig.builder().build(); + assertThrows( + IllegalStateException.class, + () -> + runner.runAsync( + "user", + sessionId, + "some-inv", + /* newMessage= */ null, + runConfig, + /* stateDelta= */ null)); + } + + // Resuming with a function response whose id matches no call in history is a caller error: + // runAsync surfaces IllegalArgumentException rather than starting a new invocation that would + // feed + // the model an orphan function response. + @Test + @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). + public void resume_functionResponseWithNoMatchingCall_throwsIllegalArgument() { + LlmAgent agent = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("x"))).name("agent").build(); + Runner runner = + Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(agent) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build()) + .build(); + Session session = runner.sessionService().createSession("test", "user").blockingGet(); + + Content orphanResponse = + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("no_such_call") + .name("pendingTool") + .response(ImmutableMap.of("status", "done"))) + .build()); + + runner + .runAsync( + "user", + session.id(), + /* invocationId= */ null, + orphanResponse, + RunConfig.builder().build(), + /* stateDelta= */ null) + .test() + .assertError(IllegalArgumentException.class); + } + + // Resuming a non-existent invocation with no new message has nothing to resume: runAsync surfaces + // IllegalArgumentException rather than running an empty model call. + @Test + @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). + public void resume_nonExistentInvocationId_throwsIllegalArgument() { + LlmAgent agent = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("x"))).name("agent").build(); + Runner runner = + Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(agent) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build()) + .build(); + Session session = runner.sessionService().createSession("test", "user").blockingGet(); + + runner + .runAsync( + "user", + session.id(), + "does-not-exist", + /* newMessage= */ null, + RunConfig.builder().build(), + /* stateDelta= */ null) + .test() + .assertError(IllegalArgumentException.class); + } + + // An orphan function response is rejected even when an explicit invocationId is supplied: the + // unmatched response is resolved first, so a provided id does not smuggle it past validation. + @Test + @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). + public void resume_orphanFunctionResponseWithProvidedInvocationId_throwsIllegalArgument() { + LlmAgent agent = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("x"))).name("agent").build(); + Runner runner = + Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(agent) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build()) + .build(); + Session session = runner.sessionService().createSession("test", "user").blockingGet(); + + Content orphanResponse = + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("no_such_call") + .name("pendingTool") + .response(ImmutableMap.of("status", "done"))) + .build()); + + runner + .runAsync( + "user", + session.id(), + "some-inv", + orphanResponse, + RunConfig.builder().build(), + /* stateDelta= */ null) + .test() + .assertError(IllegalArgumentException.class); + } + + // InMemoryRunnerTest parity: the appended function-response inherits the branch of the function + // call it answers. + @Test + @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). + public void resume_withFunctionResponse_copiesBranchFromMatchingCall() { + TestBaseAgent agent = + new TestBaseAgent("test_agent", "desc", () -> Flowable.empty(), null, null, null); + Runner runner = + Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(agent) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build()) + .build(); + Session session = runner.sessionService().createSession("test", "user").blockingGet(); + Object unusedFc = runner - .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) + .sessionService() + .appendEvent( + session, + Event.builder() + .id("fc1") + .invocationId("test-inv") + .author("test_agent") + .branch("my_special_branch") + .content( + Content.fromParts( + Part.builder() + .functionCall( + FunctionCall.builder().id("call_abc").name("test_func").build()) + .build())) + .build()) + .blockingGet(); + + Object unused = + runner + .runAsync( + "user", + session.id(), + /* invocationId= */ null, + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("call_abc") + .name("test_func") + .response(ImmutableMap.of("result", "ok"))) + .build()), + RunConfig.builder().build(), + /* stateDelta= */ null) .toList() .blockingGet(); - // Ended after the single long-running call: no function response, no second model call. - assertThat(testLlm.getRequests()).hasSize(1); - assertThat(simplifyEvents(events)).doesNotContain("agent: should not be reached"); + Session reloaded = + runner + .sessionService() + .getSession("test", "user", session.id(), Optional.empty()) + .blockingGet(); + Event lastUser = + Streams.findLast(reloaded.events().stream().filter(event -> event.author().equals("user"))) + .get(); + assertThat(lastUser.branch()).hasValue("my_special_branch"); } // The long-running call event is now a final response, but it carries no text. An agent with an @@ -2766,7 +4539,7 @@ public void runAsync_loopAgentWithLongRunningSubAgent_resumable_stopsAfterFirstI calls.incrementAndGet() <= 5 ? Flowable.just( createFunctionCallLlmResponse( - "lro_call_id", "echoTool", ImmutableMap.of("message", "hello"))) + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hello"))) : Flowable.just(createTextLlmResponse("stop"))); LlmAgent inner = createTestAgentBuilder(loopLlm) @@ -2774,7 +4547,7 @@ public void runAsync_loopAgentWithLongRunningSubAgent_resumable_stopsAfterFirstI .tools( FunctionTool.create( Tools.class, - "echoTool", + "pendingTool", /* requireConfirmation= */ false, /* isLongRunning= */ true)) .build(); @@ -2814,7 +4587,7 @@ public void runAsync_parallelAgentWithLongRunningBranch_resumable_otherBranchCom TestLlm longRunningLlm = createTestLlm( createFunctionCallLlmResponse( - "lro_call_id", "echoTool", ImmutableMap.of("message", "hello")), + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hello")), createTextLlmResponse("unexpected")); LlmAgent longRunningBranch = createTestAgentBuilder(longRunningLlm) @@ -2822,7 +4595,7 @@ public void runAsync_parallelAgentWithLongRunningBranch_resumable_otherBranchCom .tools( FunctionTool.create( Tools.class, - "echoTool", + "pendingTool", /* requireConfirmation= */ false, /* isLongRunning= */ true)) .build(); @@ -2901,7 +4674,6 @@ public void runAsync_parallelAgentWithLongRunningBranch_resumable_otherBranchCom // ResumabilityConfig is off by default and reflects the configured value. @Test - @SuppressWarnings("deprecation") // ResumabilityConfig is intentionally deprecated (partial). public void resumabilityConfig_defaultsToNotResumable() { assertThat(ResumabilityConfig.builder().build().isResumable()).isFalse(); assertThat(ResumabilityConfig.builder().resumable(true).build().isResumable()).isTrue(); diff --git a/core/src/test/java/com/google/adk/sessions/SessionJsonConverterTest.java b/core/src/test/java/com/google/adk/sessions/SessionJsonConverterTest.java index a63e3b38d..90ca6a175 100644 --- a/core/src/test/java/com/google/adk/sessions/SessionJsonConverterTest.java +++ b/core/src/test/java/com/google/adk/sessions/SessionJsonConverterTest.java @@ -276,6 +276,46 @@ public void fromApiEvent_complexActions_success() { assertThat(eventActions.endOfAgent()).isTrue(); } + @Test + public void convertEventToJson_agentState_success() throws JsonProcessingException { + EventActions actions = + EventActions.builder() + .agentState(ImmutableMap.of("current_sub_agent", "b_agent", "times_looped", 2)) + .build(); + Event event = + Event.builder() + .author("agent") + .invocationId("inv-1") + .timestamp(Instant.parse("2023-01-01T00:00:00.123Z").toEpochMilli()) + .actions(actions) + .build(); + + String json = SessionJsonConverter.convertEventToJson(event, true); + JsonNode actionsNode = objectMapper.readTree(json).get("actions"); + + assertThat(actionsNode.get("agentState").get("current_sub_agent").asText()) + .isEqualTo("b_agent"); + assertThat(actionsNode.get("agentState").get("times_looped").asInt()).isEqualTo(2); + } + + @Test + public void fromApiEvent_agentState_success() { + Map apiEvent = new HashMap<>(); + apiEvent.put("name", "sessions/123/events/456"); + apiEvent.put("invocationId", "inv-1"); + apiEvent.put("author", "agent"); + apiEvent.put("timestamp", "2023-01-01T00:00:00.123Z"); + Map actions = new HashMap<>(); + actions.put("agentState", ImmutableMap.of("current_sub_agent", "b_agent", "times_looped", 2)); + apiEvent.put("actions", actions); + + Event event = SessionJsonConverter.fromApiEvent(apiEvent); + + assertThat(event.actions().agentState()).isPresent(); + assertThat(event.actions().agentState().get()).containsEntry("current_sub_agent", "b_agent"); + assertThat(event.actions().agentState().get()).containsEntry("times_looped", 2); + } + @Test public void fromApiEvent_minimalEvent_success() { Map apiEvent = new HashMap<>(); diff --git a/core/src/test/java/com/google/adk/testing/TestUtils.java b/core/src/test/java/com/google/adk/testing/TestUtils.java index daed8d2e4..3820c8768 100644 --- a/core/src/test/java/com/google/adk/testing/TestUtils.java +++ b/core/src/test/java/com/google/adk/testing/TestUtils.java @@ -24,6 +24,7 @@ import com.google.adk.agents.InvocationContext; import com.google.adk.agents.LlmAgent; import com.google.adk.agents.RunConfig; +import com.google.adk.apps.ResumabilityConfig; import com.google.adk.artifacts.InMemoryArtifactService; import com.google.adk.events.Event; import com.google.adk.events.EventActions; @@ -71,6 +72,21 @@ public static InvocationContext createInvocationContext(BaseAgent agent) { return createInvocationContext(agent, RunConfig.builder().build()); } + /** Like {@link #createInvocationContext(BaseAgent)} but with resumability enabled. */ + public static InvocationContext createResumableInvocationContext(BaseAgent agent) { + InMemorySessionService sessionService = new InMemorySessionService(); + return InvocationContext.builder() + .sessionService(sessionService) + .artifactService(new InMemoryArtifactService()) + .invocationId("invocationId") + .agent(agent) + .session(sessionService.createSession("test_app", "test-user").blockingGet()) + .userContent(Content.fromParts(Part.fromText("user content"))) + .runConfig(RunConfig.builder().build()) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build(); + } + public static InvocationContext createInvocationContext( BaseAgent agent, BaseSessionService sessionService, Session session) { return InvocationContext.builder() @@ -105,6 +121,30 @@ public static ImmutableList simplifyEvents(List events) { .collect(toImmutableList()); } + /** Marker rendered for an end-of-agent checkpoint event by {@link #simplifyResumableEvents}. */ + public static final String END_OF_AGENT = "end_of_agent"; + + /** + * Like {@link #simplifyEvents} but renders resumability checkpoint events distinctly: an + * end-of-agent event as {@link #END_OF_AGENT} and an agent-state event as {@code + * agent_state=...}. + */ + public static ImmutableList simplifyResumableEvents(List events) { + return events.stream() + .map(event -> event.author() + ": " + formatResumableEvent(event)) + .collect(toImmutableList()); + } + + private static String formatResumableEvent(Event event) { + if (event.actions().endOfAgent()) { + return END_OF_AGENT; + } + if (event.actions().agentState().isPresent()) { + return "agent_state=" + event.actions().agentState().get(); + } + return formatEventContent(event); + } + private static String formatEventContent(Event event) { return formatContent( event