From e43b952b1e79a6636f26dafcb921e09435f01fdc Mon Sep 17 00:00:00 2001 From: Damian Momot Date: Sun, 30 Aug 2026 23:53:02 -0700 Subject: [PATCH] feat: support pausing and resuming a session invocation When enabled via ResumabilityConfig, an invocation that issues a long-running function call pauses and can be resumed later from durable per-agent checkpoints for LlmAgent, SequentialAgent, LoopAgent, and ParallelAgent, via a runAsync overload that takes an invocation id. A resumed transfer continues at the transferred-to sub-agent; a call a previous run never executed is replayed rather than re-requested from the model; and resuming persists the incoming message and any state delta before running. Behavior is unchanged unless resumability is enabled. PiperOrigin-RevId: 973694952 --- .../java/com/google/adk/agents/BaseAgent.java | 49 +- .../google/adk/agents/InvocationContext.java | 269 ++- .../java/com/google/adk/agents/LlmAgent.java | 150 +- .../java/com/google/adk/agents/LoopAgent.java | 140 +- .../com/google/adk/agents/ParallelAgent.java | 97 +- .../google/adk/agents/SequentialAgent.java | 104 +- .../adk/agents/WorkflowAgentResumption.java | 11 +- .../adk/agents/WorkflowAgentStates.java | 55 + .../main/java/com/google/adk/apps/App.java | 12 +- .../google/adk/apps/ResumabilityConfig.java | 50 +- .../adk/flows/llmflows/BaseLlmFlow.java | 136 +- .../google/adk/flows/llmflows/Functions.java | 8 +- .../google/adk/flows/llmflows/StepResume.java | 238 +++ .../java/com/google/adk/runner/Runner.java | 593 ++++++- .../adk/sessions/SessionJsonConverter.java | 9 + .../adk/agents/InvocationContextTest.java | 442 +++++ .../com/google/adk/agents/LlmAgentTest.java | 240 +++ .../com/google/adk/agents/LoopAgentTest.java | 123 ++ .../google/adk/agents/ParallelAgentTest.java | 202 +++ .../adk/agents/SequentialAgentTest.java | 146 ++ .../adk/apps/ResumabilityConfigTest.java | 65 + .../adk/flows/llmflows/FunctionsTest.java | 36 + .../runner/RunnerLegacyResumabilityTest.java | 728 ++++++++ .../adk/runner/RunnerResumabilityTest.java | 1514 +++++++++++++++++ .../com/google/adk/runner/RunnerTest.java | 327 +--- .../sessions/SessionJsonConverterTest.java | 58 + .../adk/testing/ResumabilityTestUtils.java | 294 ++++ .../com/google/adk/testing/TestUtils.java | 40 + 28 files changed, 5663 insertions(+), 473 deletions(-) create mode 100644 core/src/main/java/com/google/adk/agents/WorkflowAgentStates.java create mode 100644 core/src/main/java/com/google/adk/flows/llmflows/StepResume.java create mode 100644 core/src/test/java/com/google/adk/apps/ResumabilityConfigTest.java create mode 100644 core/src/test/java/com/google/adk/runner/RunnerLegacyResumabilityTest.java create mode 100644 core/src/test/java/com/google/adk/runner/RunnerResumabilityTest.java create mode 100644 core/src/test/java/com/google/adk/testing/ResumabilityTestUtils.java 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..87191c3b5 100644 --- a/core/src/main/java/com/google/adk/agents/InvocationContext.java +++ b/core/src/main/java/com/google/adk/agents/InvocationContext.java @@ -18,8 +18,10 @@ import static com.google.common.base.Strings.isNullOrEmpty; +import com.google.adk.annotations.Experimental; 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 +29,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 +68,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 +99,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 +240,249 @@ 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. */ + @Experimental public boolean isResumable() { return resumabilityConfig != null && resumabilityConfig.isResumable(); } + /** + * Returns whether the invocation runs the legacy resumption flow: the behavior {@link + * #isResumable()} had before durable checkpoints existed. It is a separate flow, not a weaker + * {@link #isResumable()}, so a caller wanting either has to ask for both. + * + * @deprecated Reports the deprecated plain-text continuation shim, and goes away with it. Use + * {@link #isResumable()}. + */ + @Deprecated + @SuppressWarnings("deprecation") // The shim it reads is deprecated by design. + public boolean isLegacyResumability() { + return resumabilityConfig != null && resumabilityConfig.isPlainTextContinuationAutoResume(); + } + + /** + * 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}. + */ + @Experimental + 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. + */ + @Experimental + 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. + */ + @Experimental + 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) + // An after-agent callback emits content after end-of-agent, which must not reopen it. + && !endOfAgents.getOrDefault(author, false)) { + 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). + */ + @Experimental + public 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 invocation is {@link #isResumable()} + * and the event carries a long-running function call (including a synthetic {@code + * adk_request_confirmation} HITL request). + */ + @Experimental + public boolean shouldPauseInvocation(Event event) { + return isResumable() && carriesLongRunningCall(event); + } + + /** Returns whether the event carries a long-running function call, independent of the mode. */ + private static boolean carriesLongRunningCall(Event event) { + 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() { + return hasUnansweredLongRunningCall( + events(/* currentInvocation= */ true, /* currentBranch= */ true), /* scope= */ null); + } + + /** + * Returns whether a long-running call made anywhere inside {@code agent}'s subtree is still + * unanswered. A workflow agent needs this because a sub-agent can pause silently -- emitting no + * event, so no event carries the pause -- and because Java gives every sub-agent of a {@link + * ParallelAgent} that agent's own branch, which no branch filter can tell apart. Scoping by + * subtree keeps a paused parallel sibling from stalling an unrelated branch, as Python does. + */ + boolean hasUnansweredLongRunningCallIn(BaseAgent agent) { + Set scope = new HashSet<>(); + collectAgentNames(agent, scope); + return hasUnansweredLongRunningCall( + events(/* currentInvocation= */ true, /* currentBranch= */ false), scope); + } + + /** Adds {@code agent}'s name and every descendant's name to {@code names}. */ + private static void collectAgentNames(BaseAgent agent, Set names) { + if (!names.add(agent.name())) { + return; + } + for (BaseAgent subAgent : agent.subAgents()) { + collectAgentNames(subAgent, names); + } + } + + /** + * Whether {@code events} hold a long-running call with no response. {@code scope}, when non-null, + * restricts which authors' calls count; responses always count whoever authored them, since a + * resumed answer is authored by the user. + */ + private static boolean hasUnansweredLongRunningCall( + ImmutableList events, @Nullable Set scope) { + if (events.isEmpty()) { + return false; + } + Set awaited = new HashSet<>(); + for (Event event : events) { + if (!carriesLongRunningCall(event)) { + continue; + } + if (scope != null && !scope.contains(event.author())) { + 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(); + } + ImmutableList 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 +542,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 +566,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..bd805afcc 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; @@ -80,6 +82,9 @@ public class LlmAgent extends BaseAgent { private static final Logger logger = LoggerFactory.getLogger(LlmAgent.class); + /** Name of the built-in tool whose response records an agent transfer. */ + private static final String TRANSFER_TO_AGENT_TOOL = "transfer_to_agent"; + /** * Enum to define if contents of previous events should be included in requests to the underlying * LLM. @@ -660,7 +665,150 @@ 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 (transferTargetNameFrom(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 name of the agent this agent transferred to in {@code event}, or empty when {@code + * event} is not such a transfer. A {@code transfer_to_agent} function response is required, as in + * Python ADK, so a callback that sets the action on an unrelated event does not read as one. + */ + private Optional transferTargetNameFrom(Event event) { + if (!name().equals(event.author()) + || event.functionResponses().stream() + .noneMatch( + response -> response.name().filter(TRANSFER_TO_AGENT_TOOL::equals).isPresent())) { + return Optional.empty(); + } + return event.actions().transferToAgent().filter(target -> !target.equals(name())); + } + + /** + * Resolves a transfer target recorded in history back to its agent. A live transfer fails in + * {@code BaseLlmFlow} when the target is gone; a resumed one has no such check, and returning + * empty here would silently re-run this agent instead of the sub-agent it handed off to. + */ + private BaseAgent resolveTransferTarget(String targetName) { + return rootAgent() + .findAgent(targetName) + .orElseThrow( + () -> + new IllegalStateException( + "Cannot resume agent " + + name() + + ": the agent it transferred to is no longer in the agent tree.")); + } + + /** + * 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) { + ImmutableList events = + context.events(/* currentInvocation= */ true, /* currentBranch= */ true); + if (events.isEmpty()) { + return Optional.empty(); + } + Event lastEvent = Iterables.getLast(events); + if (name().equals(lastEvent.author())) { + return transferTargetNameFrom(lastEvent).map(this::resolveTransferTarget); + } + 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 targetName = transferTargetNameFrom(events.get(i)); + if (targetName.isPresent()) { + return Optional.of(resolveTransferTarget(targetName.get())); + } + } + 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..f5155c52d 100644 --- a/core/src/main/java/com/google/adk/agents/LoopAgent.java +++ b/core/src/main/java/com/google/adk/agents/LoopAgent.java @@ -16,11 +16,15 @@ package com.google.adk.agents; +import static com.google.common.base.Strings.isNullOrEmpty; + 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; @@ -136,12 +140,17 @@ public static LoopAgent fromConfig(LoopAgentConfig config, String configAbsPath) } @Override + @SuppressWarnings("deprecation") // The shim it dispatches on is deprecated by design. protected Flowable runAsyncImpl(InvocationContext invocationContext) { List subAgents = subAgents(); if (subAgents == null || subAgents.isEmpty()) { return Flowable.empty(); } + if (invocationContext.isLegacyResumability()) { + return runAsyncLegacyResumption(invocationContext, subAgents); + } + if (!invocationContext.isResumable()) { return Flowable.fromIterable(subAgents) .concatMap(subAgent -> subAgent.runAsync(invocationContext)) @@ -149,16 +158,46 @@ 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). + // 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; + // Empty counts as absent, as in Python: start at the first sub-agent rather than warning + // about a sub-agent that was never named. + int startIndex = + isNullOrEmpty(startSubAgentName) + ? 0 + : WorkflowAgentStates.findIndexForResumption( + subAgents, startSubAgentName, logger); + LoopState loopState = new LoopState(state != null, startTimesLooped); + return runLoopIteration(invocationContext, subAgents, startIndex, loopState); + }); + } + + /** + * Runs the loop under the deprecated legacy resumption flow. Frozen copy of the behavior + * resumability had before durable checkpoints: stop looping once a sub-agent emits a pending + * long-running call, with no state read or written, so a resume cannot re-enter that iteration. + */ + private Flowable runAsyncLegacyResumption( + InvocationContext invocationContext, List subAgents) { AtomicBoolean paused = new AtomicBoolean(false); AtomicInteger timesLooped = new AtomicInteger(0); return Flowable.fromIterable(subAgents) .concatMap( subAgent -> paused.get() - ? Flowable.empty() + ? Flowable.empty() : subAgent .runAsync(invocationContext) .doOnNext( @@ -174,6 +213,99 @@ protected Flowable runAsyncImpl(InvocationContext invocationContext) { .takeUntil(LoopAgent::hasEscalateAction); } + /** 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(); + } + // The previous sub-agent may have paused silently -- emitting + // nothing, so no event carried the pause -- and under a + // ParallelAgent its call sits on a branch this agent cannot see. + if (context.hasUnansweredLongRunningCallIn(this)) { + state.paused.set(true); + 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() || context.hasUnansweredLongRunningCallIn(this)) { + return Flowable.empty(); + } + // Anything but a pause completes the iteration, so count it and clear the + // sub-agents' state before exiting or looping again. + state.timesLooped.incrementAndGet(); + context.resetSubAgentStates(name()); + if (state.shouldExit.get()) { + return endOfAgentAndRecord(context); + } + // 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 protected Flowable runLiveImpl(InvocationContext invocationContext) { return Flowable.error( 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..6c0df73da 100644 --- a/core/src/main/java/com/google/adk/agents/ParallelAgent.java +++ b/core/src/main/java/com/google/adk/agents/ParallelAgent.java @@ -16,15 +16,19 @@ package com.google.adk.agents; import static com.google.common.base.Strings.isNullOrEmpty; +import static com.google.common.collect.ImmutableSet.toImmutableSet; import com.google.adk.agents.ConfigAgentUtils.ConfigurationException; import com.google.adk.events.Event; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; 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 +179,94 @@ 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)); + } + // Unscoped on purpose: any escalate stops the branches, as it did before resumability. + // Narrowing it here would change the default path, not just this CL's new one; the resumable + // path below scopes it to direct sub-agents, as Python does. + return Flowable.merge(agentFlowables) + .takeUntil((Event event) -> event.actions().escalate().orElse(false)); } - return Flowable.merge(agentFlowables) - .takeUntil((Event event) -> event.actions().escalate().orElse(false)); + + ImmutableSet subAgentNames = + currentSubAgents.stream().map(BaseAgent::name).collect(toImmutableSet()); + + // 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 (asksThisAgentToExit(event, subAgentNames)) { + escalated.set(true); + } + })); + } + Flowable merged = + Flowable.merge(agentFlowables) + .takeUntil((Event event) -> asksThisAgentToExit(event, subAgentNames)); + + 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(); + })); + }); + } + + /** + * Returns whether {@code event} asks this agent to stop its remaining branches. An escalation + * ends the workflow that directly encloses the escalating agent, and that workflow re-yields the + * event while unwinding, so only one authored by a direct sub-agent is addressed to this agent. + */ + private static boolean asksThisAgentToExit(Event event, ImmutableSet subAgentNames) { + return event.actions().escalate().orElse(false) && subAgentNames.contains(event.author()); } /** 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..f4148a842 100644 --- a/core/src/main/java/com/google/adk/agents/SequentialAgent.java +++ b/core/src/main/java/com/google/adk/agents/SequentialAgent.java @@ -15,10 +15,14 @@ */ package com.google.adk.agents; +import static com.google.common.base.Strings.isNullOrEmpty; + 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.concurrent.atomic.AtomicBoolean; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -90,24 +94,38 @@ public static Builder builder() { /** * Runs sub-agents sequentially. * - *

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. + *

Three modes: with resumability on, a resume fast-forwards to the checkpointed sub-agent and + * pauses on a long-running call; with the deprecated shim on, the same happens but the resume + * point is reconstructed from history and nothing is checkpointed; with neither, sub-agents just + * run in order. * * @param invocationContext Invocation context. * @return Flowable emitting events from sub-agents. */ @Override + @SuppressWarnings("deprecation") // The shim it dispatches on is deprecated by design. protected Flowable runAsyncImpl(InvocationContext invocationContext) { List subAgents = subAgents(); if (subAgents.isEmpty()) { return Flowable.empty(); } - if (!invocationContext.isResumable()) { - return Flowable.fromIterable(subAgents) - .concatMap(subAgent -> subAgent.runAsync(invocationContext)); + if (invocationContext.isResumable()) { + return runAsyncResumable(invocationContext, subAgents); + } + if (invocationContext.isLegacyResumability()) { + return runAsyncLegacyResumption(invocationContext, subAgents); } + return Flowable.fromIterable(subAgents) + .concatMap(subAgent -> subAgent.runAsync(invocationContext)); + } + + /** + * Runs sub-agents under the deprecated legacy resumption flow, reconstructing the resume point + * from session events. Frozen copy of the behavior resumability had before durable checkpoints: + * no state is read or written, so only history decides where the sequence restarts. + */ + private Flowable runAsyncLegacyResumption( + InvocationContext invocationContext, List subAgents) { int startIndex = WorkflowAgentResumption.resumeSubAgentIndex(invocationContext, subAgents).orElse(0); AtomicBoolean paused = new AtomicBoolean(false); @@ -126,6 +144,78 @@ protected Flowable runAsyncImpl(InvocationContext invocationContext) { })); } + /** + * Runs sub-agents under durable resumability, matching Python ADK: checkpoint each sub-agent + * before it runs, fast-forward to the checkpoint on resume, and pause (without ending) on a + * long-running call. + */ + private Flowable runAsyncResumable( + InvocationContext invocationContext, List subAgents) { + // Deferred so the checkpoint is read and the loop's mutable state created per subscription, + // not once when the Flowable is assembled. + 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 startIndex; + if (state == null) { + startIndex = 0; + } else if (isNullOrEmpty(startSubAgentName)) { + // State with no current sub-agent means the sequence already finished; empty counts as + // absent, since that is Python's default for the field. + startIndex = subAgents.size(); + } else { + startIndex = + WorkflowAgentStates.findIndexForResumption(subAgents, startSubAgentName, logger); + } + AtomicBoolean paused = new AtomicBoolean(false); + AtomicBoolean resuming = new AtomicBoolean(state != null); + return Flowable.fromIterable(subAgents.subList(startIndex, subAgents.size())) + .concatMap( + subAgent -> + Flowable.defer( + () -> { + // The previous sub-agent may have paused silently, emitting nothing, so + // no event carried the pause. Scoped to this agent's subtree so a + // paused parallel sibling elsewhere does not stall this sequence. + if (paused.get() + || invocationContext.hasUnansweredLongRunningCallIn(this)) { + paused.set(true); + 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() + || invocationContext.hasUnansweredLongRunningCallIn(this)) { + return Flowable.empty(); + } + return endOfAgentAndRecord(invocationContext); + })); + }); + } + /** * Runs sub-agents sequentially in live mode. * 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..811be0c89 100644 --- a/core/src/main/java/com/google/adk/agents/WorkflowAgentResumption.java +++ b/core/src/main/java/com/google/adk/agents/WorkflowAgentResumption.java @@ -22,14 +22,15 @@ import java.util.Optional; /** - * Helpers for resuming workflow agents from session events. Temporary until session resumption - * (persisted agent state) is available. + * Legacy-resumption helper: reconstructs a workflow agent's resume point from session events. Used + * only by the legacy flow, which writes no durable agent state; the resumable flow resumes from its + * checkpoints instead, as Python ADK does. */ final class WorkflowAgentResumption { /** - * Index of the direct sub-agent whose subtree authored the call the latest event resumes, or - * empty when not resuming into this workflow. + * Index of the direct sub-agent whose subtree authored the call being resumed, or empty when not + * resuming into this workflow. */ static Optional resumeSubAgentIndex( InvocationContext invocationContext, List subAgents) { @@ -50,6 +51,8 @@ static Optional resumeSubAgentIndex( /** * Whether the event emits a long-running call still awaiting a response (e.g. a HITL request). + * The legacy flow pauses on this rather than on {@link InvocationContext#shouldPauseInvocation}, + * which the resumable flow uses. */ static boolean hasPendingLongRunningCall(Event event) { return Functions.hasPendingLongRunningCall(event); 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..2ea9b30ff 100644 --- a/core/src/main/java/com/google/adk/apps/ResumabilityConfig.java +++ b/core/src/main/java/com/google/adk/apps/ResumabilityConfig.java @@ -16,34 +16,43 @@ package com.google.adk.apps; +import static com.google.common.base.Preconditions.checkArgument; + +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. + * + *

The two flags select the resumption flow and are mutually exclusive. With neither set the app + * never pauses; {@link #isResumable()} checkpoints agent state and resumes from it; {@link + * #isPlainTextContinuationAutoResume()} selects the deprecated legacy flow. * - * @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 { - /** Whether the app supports agent resumption. */ + /** Whether the app supports agent resumption, checkpointing agent state as it runs. */ public abstract boolean isResumable(); /** * 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. + * + *

Selects the legacy resumption flow, which reconstructs the resume point from session events + * and never persists agent state. Mutually exclusive with {@link #isResumable()}. * * @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)} - * (or send a function response to the paused call) and stop setting this flag; it will be + * {@code Runner.runAsync(userId, sessionId, invocationId, message, runConfig, stateDelta)}, + * which resumes the invocation named by {@code invocationId}, or answer the paused call with + * a function response; both need {@link #isResumable()} instead of this flag, which will be * removed. */ @Deprecated @@ -70,6 +79,21 @@ public abstract static class Builder { @CanIgnoreReturnValue public abstract Builder plainTextContinuationAutoResume(boolean value); - public abstract ResumabilityConfig build(); + abstract ResumabilityConfig autoBuild(); + + /** + * Builds the config, rejecting a combination of flags that has no defined behavior. + * + * @throws IllegalArgumentException if both resumability and the legacy shim are enabled; they + * select different resumption flows, so exactly one may be set. + */ + public ResumabilityConfig build() { + ResumabilityConfig config = autoBuild(); + checkArgument( + !(config.isResumable() && config.isPlainTextContinuationAutoResume()), + "resumable and plainTextContinuationAutoResume are mutually exclusive: set resumable for" + + " the supported flow, or the deprecated shim for the legacy one."); + return config; + } } } 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..8f3114204 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 @@ -60,6 +60,7 @@ import java.util.List; import java.util.Optional; import java.util.Set; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; import java.util.function.BiFunction; import org.slf4j.Logger; @@ -432,7 +433,14 @@ private Flowable runOneStep(Context spanContext, InvocationContext contex return Flowable.defer( () -> { + AtomicBoolean preprocessAnswered = new AtomicBoolean(false); return preprocess(context, llmRequestRef) + .doOnNext( + event -> { + if (!event.functionResponses().isEmpty() && event.finalResponse()) { + preprocessAnswered.set(true); + } + }) .concatWith( Flowable.defer( () -> { @@ -441,6 +449,33 @@ private Flowable runOneStep(Context spanContext, InvocationContext contex logger.debug("End invocation requested during preprocessing."); return Flowable.empty(); } + // A tool confirmed and executed during preprocessing already answered this + // step, so the resume decision below would replay it. Gated on + // resumability, since only the resumable flow makes that decision. + if (context.isResumable() && preprocessAnswered.get()) { + logger.debug("Preprocessing produced a final response; ending the step."); + return Flowable.empty(); + } + + // Decide before calling the model, as Python does: a resumed branch may + // still owe an answer, or owe a call a previous run never executed. + StepResume.Decision resume = + StepResume.decide(context, llmRequestAfterPreprocess.tools()); + if (resume.action == StepResume.Action.PAUSE) { + logger.debug("Pausing the flow: a call is still unanswered."); + return Flowable.empty(); + } + if (resume.action == StepResume.Action.REPLAY_CALLS) { + logger.debug("Replaying function calls a previous run did not execute."); + // Same follow-ups as a fresh call, including a replayed transfer: the + // calls were persisted but never executed, so nothing downstream ran. + return runFunctionCalls( + context, + resume.replayEvent(), + llmRequestAfterPreprocess, + spanContext) + .concatMap(event -> followTransfer(event, context, spanContext)); + } try { context.incrementLlmCallsCount(); @@ -477,27 +512,7 @@ private Flowable runOneStep(Context spanContext, InvocationContext contex String newId = Event.generateEventId(); logger.debug("Resetting event ID from {} to {}", oldId, newId); event = event.toBuilder().id(newId).build(); - Flowable postProcessedEvents = Flowable.just(event); - if (event.actions().transferToAgent().isPresent()) { - String agentToTransfer = - event.actions().transferToAgent().get(); - BaseAgent rootAgent = context.agent().rootAgent(); - Optional nextAgent = - rootAgent.findAgent(agentToTransfer); - if (nextAgent.isEmpty()) { - logger.error("Agent not found: {}", agentToTransfer); - return postProcessedEvents.concatWith( - Flowable.error( - new IllegalStateException( - "Agent not found: " + agentToTransfer))); - } - return postProcessedEvents.concatWith( - nextAgent - .get() - .runAsync(context) - .compose(Tracing.withContext(spanContext))); - } - return postProcessedEvents; + return followTransfer(event, context, spanContext); }); })); }); @@ -521,6 +536,8 @@ private Flowable run( logger.debug("Ending flow execution because max steps reached."); return currentStepEvents; } + @SuppressWarnings("deprecation") // The shim it reports is deprecated by design. + boolean legacyResumption = invocationContext.isLegacyResumability(); return currentStepEvents.concatWith( currentStepEvents @@ -534,11 +551,9 @@ private Flowable run( "Ending flow execution based on final response, endInvocation action or" + " empty event list."); 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. + } else if (legacyResumption && Functions.hasPendingLongRunningCall(eventList)) { + // Legacy resumption pauses on an unanswered long-running call here; the + // resumable flow decides before the model call instead, in StepResume. logger.debug("Pausing flow execution on a pending long-running call."); return Flowable.empty(); } else { @@ -758,30 +773,59 @@ private Flowable buildPostprocessingEvents( return processorEvents.concatWith(Flowable.just(modelResponseEvent)); } - Flowable functionEvents; + Flowable functionEvents = + runFunctionCalls(context, modelResponseEvent, llmRequest, parentContext); + + return processorEvents.concatWith(Flowable.just(modelResponseEvent)).concatWith(functionEvents); + } + + /** + * Emits {@code event}, followed by the transferred-to agent's run when it carries a transfer. + * Shared so a replayed transfer continues into its target the way a fresh one does. + */ + private Flowable followTransfer( + Event event, InvocationContext context, Context spanContext) { + Flowable self = Flowable.just(event); + if (event.actions().transferToAgent().isEmpty()) { + return self; + } + String agentToTransfer = event.actions().transferToAgent().get(); + Optional nextAgent = context.agent().rootAgent().findAgent(agentToTransfer); + if (nextAgent.isEmpty()) { + logger.error("Agent not found: {}", agentToTransfer); + return self.concatWith( + Flowable.error(new IllegalStateException("Agent not found: " + agentToTransfer))); + } + return self.concatWith( + nextAgent.get().runAsync(context).compose(Tracing.withContext(spanContext))); + } + + /** + * Runs {@code callEvent}'s function calls and emits what follows from them: the tool-confirmation + * request, the function response, and any structured final response. Shared by the normal + * postprocessing path and by replaying calls a previous run never executed. + */ + private Flowable runFunctionCalls( + InvocationContext context, Event callEvent, LlmRequest llmRequest, Context parentContext) { try (Scope scope = parentContext.makeCurrent()) { Maybe maybeFunctionResponseEvent = context.runConfig().streamingMode() == StreamingMode.BIDI - ? Functions.handleFunctionCallsLive(context, modelResponseEvent, llmRequest.tools()) - : Functions.handleFunctionCalls(context, modelResponseEvent, llmRequest.tools()); - functionEvents = - maybeFunctionResponseEvent.flatMapPublisher( - functionResponseEvent -> { - Optional toolConfirmationEvent = - Functions.generateRequestConfirmationEvent( - context, modelResponseEvent, functionResponseEvent); - List events = new ArrayList<>(); - toolConfirmationEvent.ifPresent(events::add); - events.add(functionResponseEvent); - OutputSchema.getStructuredModelResponse(functionResponseEvent) - .ifPresent( - json -> - events.add(OutputSchema.createFinalModelResponseEvent(context, json))); - return Flowable.fromIterable(events); - }); + ? Functions.handleFunctionCallsLive(context, callEvent, llmRequest.tools()) + : Functions.handleFunctionCalls(context, callEvent, llmRequest.tools()); + return maybeFunctionResponseEvent.flatMapPublisher( + functionResponseEvent -> { + Optional toolConfirmationEvent = + Functions.generateRequestConfirmationEvent( + context, callEvent, functionResponseEvent); + List events = new ArrayList<>(); + toolConfirmationEvent.ifPresent(events::add); + events.add(functionResponseEvent); + OutputSchema.getStructuredModelResponse(functionResponseEvent) + .ifPresent( + json -> events.add(OutputSchema.createFinalModelResponseEvent(context, json))); + return Flowable.fromIterable(events); + }); } - - return processorEvents.concatWith(Flowable.just(modelResponseEvent)).concatWith(functionEvents); } /** 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..cd0528743 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,9 +472,11 @@ 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 either of the last two events emits a long-running call, meaning the legacy + * resumption flow should pause instead of calling the model again. Responses are not matched + * against calls: a long-running tool that returns a value in the same turn still pauses, because + * its call event is inside the window. The resumable flow does not use this -- it decides in + * {@code StepResume}, which does match responses. */ static boolean hasPendingLongRunningCall(List events) { int from = Math.max(0, events.size() - 2); diff --git a/core/src/main/java/com/google/adk/flows/llmflows/StepResume.java b/core/src/main/java/com/google/adk/flows/llmflows/StepResume.java new file mode 100644 index 000000000..7cfd16e98 --- /dev/null +++ b/core/src/main/java/com/google/adk/flows/llmflows/StepResume.java @@ -0,0 +1,238 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.adk.flows.llmflows; + +import com.google.adk.agents.InvocationContext; +import com.google.adk.events.Event; +import com.google.adk.tools.BaseTool; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; +import com.google.genai.types.FunctionCall; +import com.google.genai.types.FunctionResponse; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.jspecify.annotations.Nullable; + +/** + * Decides how a resumable flow's next step continues: carry on to the model, stop because a call is + * still unanswered, or run calls a previous run never executed. + * + *

Ports Python ADK's resume decision. Java has no per-call branch run ids, so the sub-branch + * case Python recognizes (a HITL answer returning against a branch the call opened) has no + * counterpart here and is treated as absent. + */ +final class StepResume { + + /** What the flow should do with the events it resumed from. */ + enum Action { + /** Nothing outstanding; proceed to the model call. */ + CONTINUE, + /** A call is still unanswered; stop without emitting anything. */ + PAUSE, + /** A call was never executed; run the calls on {@link Decision#event}. */ + REPLAY_CALLS + } + + /** The action to take, and the event it applies to. */ + static final class Decision { + final Action action; + private final @Nullable Event event; + + private Decision(Action action, @Nullable Event event) { + this.action = action; + this.event = event; + } + + /** The event whose calls to run; only a {@code REPLAY_CALLS} decision carries one. */ + Event replayEvent() { + if (event == null) { + throw new IllegalStateException(action + " decision carries no event to replay"); + } + return event; + } + } + + private static final Decision CONTINUE = new Decision(Action.CONTINUE, null); + private static final Decision PAUSE = new Decision(Action.PAUSE, null); + + /** + * Decides how the next step resumes. The branch's last event decides the case: content means the + * normal flow carries on, while a call that was never executed has to run first. + */ + static Decision decide(InvocationContext context, Map tools) { + if (!context.isResumable()) { + return CONTINUE; + } + ImmutableList events = + context.events(/* currentInvocation= */ true, /* currentBranch= */ true); + if (events.isEmpty()) { + return CONTINUE; + } + if (events.size() > 1) { + Decision decision = decideFromBranch(context, events, tools); + if (decision.action != Action.CONTINUE) { + return decision; + } + } + // A single event, or a branch the scan cleared: the branch still owes something only when its + // last event carries calls nothing has answered -- being last is what makes them unanswered. + Event last = events.get(events.size() - 1); + if (!last.partial().orElse(false) + && !last.functionCalls().isEmpty() + // Only replay a call the CURRENT AGENT authored: a restored session may carry + // caller-supplied events, and dispatching one would run a tool with no model turn. + && context.agent().name().equals(last.author())) { + return new Decision(Action.REPLAY_CALLS, last); + } + return CONTINUE; + } + + /** The multi-event core: pause on an unanswered call, else replay an unexecuted one. */ + private static Decision decideFromBranch( + InvocationContext context, List events, Map tools) { + Event last = events.get(events.size() - 1); + boolean pausedByLast = context.shouldPauseInvocation(last); + if (!pausedByLast && pauseLeftCallsUnanswered(context, events)) { + return PAUSE; + } + + int callIdx = findTargetCallEventIndex(context, events, tools); + if (callIdx < 0) { + return pausedByLast ? PAUSE : CONTINUE; + } + Event callEvent = events.get(callIdx); + Set callNames = new HashSet<>(); + Set callIds = new HashSet<>(); + for (FunctionCall call : callEvent.functionCalls()) { + call.name().ifPresent(callNames::add); + call.id().ifPresent(callIds::add); + } + Set longRunningIds = new HashSet<>(); + for (int i = callIdx; i < events.size(); i++) { + longRunningIds.addAll(events.get(i).longRunningToolIds().orElse(ImmutableSet.of())); + } + callIds.addAll(longRunningIds); + + Set answeredIds = new HashSet<>(); + for (int i = callIdx + 1; i < events.size(); i++) { + for (FunctionResponse response : events.get(i).functionResponses()) { + response.id().ifPresent(answeredIds::add); + } + } + ImmutableList answers = answerEvent(events, callIdx, callIds, callNames); + + boolean lroUnanswered = !longRunningIds.isEmpty() && disjoint(longRunningIds, answeredIds); + boolean callUnanswered = + !callIds.isEmpty() + && disjoint(callIds, answeredIds) + && answers.stream() + .noneMatch(response -> response.name().map(callNames::contains).orElse(false)); + if (lroUnanswered || callUnanswered) { + return PAUSE; + } + if (needsCallReplay(callNames, answers)) { + return new Decision(Action.REPLAY_CALLS, callEvent); + } + return pausedByLast ? PAUSE : CONTINUE; + } + + /** + * Whether a pause earlier in {@code events} is still waiting. Every event before the last counts: + * a long-running call followed by several responses leaves the pausing call further back than a + * two-event window can see. + */ + private static boolean pauseLeftCallsUnanswered(InvocationContext context, List events) { + Set awaited = new HashSet<>(); + for (int i = 0; i < events.size() - 1; i++) { + Event event = events.get(i); + if (!context.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); + } + } + // Asks whether ANY awaited id is still open, so a partially answered pause keeps waiting. + return !answered.containsAll(awaited); + } + + /** + * Index of the most recent event before the last that calls a tool this flow owns and the current + * agent authored, or -1 when there is none. + */ + private static int findTargetCallEventIndex( + InvocationContext context, List events, Map tools) { + String agentName = context.agent().name(); + for (int i = events.size() - 2; i >= 0; i--) { + Event event = events.get(i); + if (!agentName.equals(event.author())) { + continue; + } + boolean callsOwnedTool = + event.functionCalls().stream() + .anyMatch(call -> call.name().map(tools::containsKey).orElse(false)); + if (callsOwnedTool) { + return i; + } + } + return -1; + } + + /** The responses of the event answering the call, or of the last event when none does. */ + private static ImmutableList answerEvent( + List events, int callIdx, Set callIds, Set callNames) { + for (int i = events.size() - 1; i > callIdx; i--) { + for (FunctionResponse response : events.get(i).functionResponses()) { + boolean matchesId = response.id().map(callIds::contains).orElse(false); + boolean matchesName = + response.id().isEmpty() && response.name().map(callNames::contains).orElse(false); + if (matchesId || matchesName) { + return events.get(i).functionResponses(); + } + } + } + return events.get(events.size() - 1).functionResponses(); + } + + /** Whether the calls still need running: nothing answered them, or something else did. */ + private static boolean needsCallReplay(Set callNames, List answers) { + if (callNames.isEmpty()) { + return false; + } + return answers.isEmpty() + || answers.stream() + .anyMatch(response -> !response.name().map(callNames::contains).orElse(false)); + } + + private static boolean disjoint(Set left, Set right) { + return left.stream().noneMatch(right::contains); + } + + private StepResume() {} +} 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 fbecbf02b..3eb77a20c 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,9 @@ package com.google.adk.runner; import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Preconditions.checkState; +import static com.google.common.collect.ImmutableList.toImmutableList; +import static java.util.stream.Collectors.toCollection; import com.google.adk.agents.ActiveStreamingTool; import com.google.adk.agents.BaseAgent; @@ -24,9 +27,12 @@ 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.ParallelAgent; 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; @@ -56,6 +62,8 @@ 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 +77,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.Objects; 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. public class Runner { private final BaseAgent agent; private final String appName; @@ -360,6 +370,22 @@ private Single createUserMessageEvent( 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 +419,7 @@ private Single createUserMessageEvent( .id(Event.generateEventId()) .invocationId(invocationContext.invocationId()) .author(Role.USER) + .branch(branch) .content(messageToAppend); // Add state delta if provided @@ -504,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. * @@ -521,6 +594,19 @@ protected Flowable runAsyncImpl( Preconditions.checkNotNull(session, "session cannot be null"); Preconditions.checkNotNull(newMessage, "newMessage cannot be null"); Preconditions.checkNotNull(runConfig, "runConfig cannot be null"); + if (isResumable()) { + 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(); @@ -558,12 +644,7 @@ protected Flowable runAsyncImpl( userEvent -> runAgentForUserEvent(initialContext, session, userEvent, rootAgent) .compose(Tracing.withContext(capturedContext))) - .doOnError( - throwable -> - this.pluginManager - .runOnRunErrorCallback(initialContext, throwable) - .onErrorComplete() - .subscribe()); + .doOnError(throwable -> runOnRunError(initialContext, throwable)); }) .doOnError( throwable -> { @@ -621,27 +702,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 @@ -651,39 +741,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) { @@ -694,6 +777,353 @@ 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( + () -> { + // Reject a message no invocation can be resumed from before any invocation-id fallback + // or auto-resume. Flowable.defer turns the thrown IAE into an error signal. + if (newMessage != null) { + validateResumeMessage(session, newMessage); + } + String resolvedInvocationId = + resolveInvocationId(session, newMessage, providedInvocationId); + 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 + ? newMessage + : originalUserMessage(session, resolvedInvocationId)) + .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)) + // Persist before running: rehydration and agent resolution read session + // history, so an unpersisted response leaves the call looking unanswered. + .flatMap(userEvent -> this.sessionService.appendEvent(session, userEvent)) + .flatMapPublisher( + userEvent -> + runResumedAgent( + session, + resolvedInvocationId, + userEvent.content().orElse(null), + runConfig)); + } else if (stateDelta != null && !stateDelta.isEmpty()) { + // No message to carry the delta, so persist it as a content-less event rather than + // leaving it in memory only, matching Python ADK. + Event stateDeltaEvent = + Event.builder() + .id(Event.generateEventId()) + .invocationId(resolvedInvocationId) + .author(Role.USER) + .actions( + EventActions.builder() + .stateDelta(new ConcurrentHashMap<>(stateDelta)) + .build()) + .build(); + events = + this.sessionService + .appendEvent(session, stateDeltaEvent) + .ignoreElement() + .andThen( + runResumedAgent( + session, resolvedInvocationId, /* userContent= */ 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 + ? userContent + : originalUserMessage(session, resolvedInvocationId)) + .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 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 = functionResponseIds(newMessage); + 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(); + } + + /** + * Rejects a resume message no invocation can be resumed from, matching Python ADK: one mixing + * text with function responses, one whose function response carries no id, one whose response ids + * match no call in session history (it would feed the model an orphan response), and one whose + * responses span more than one invocation. Counts rather than ids are reported, since the ids + * come from the caller. + * + * @throws IllegalArgumentException if {@code newMessage} is not resumable. + */ + private static void validateResumeMessage(Session session, Content newMessage) { + ImmutableList responses = + newMessage.parts().stream() + .flatMap(List::stream) + .map(Part::functionResponse) + .flatMap(Optional::stream) + .collect(toImmutableList()); + if (responses.isEmpty()) { + return; + } + checkArgument( + newMessage.parts().stream() + .flatMap(List::stream) + .noneMatch(part -> part.text().isPresent()), + "A resume message cannot carry both function responses and text: function responses resume" + + " an existing invocation while text starts a new one."); + checkArgument( + responses.stream().allMatch(response -> response.id().isPresent()), + "A function response id is required to resume an invocation."); + Set unmatched = functionResponseIds(newMessage); + Set invocationIds = new HashSet<>(); + List events = session.events(); + for (int i = events.size() - 1; i >= 0 && !unmatched.isEmpty(); i--) { + Event event = events.get(i); + for (FunctionCall call : event.functionCalls()) { + if (call.id().filter(unmatched::remove).isPresent()) { + invocationIds.add(event.invocationId()); + } + } + } + checkArgument( + unmatched.isEmpty(), + "No matching function call for %s of the resume message's function responses.", + unmatched.size()); + checkArgument( + invocationIds.size() <= 1, + "The resume message's function responses span %s invocations; all of them must answer the" + + " same one.", + invocationIds.size()); + } + + /** Returns the ids of every function response part {@code message} carries. */ + private static Set functionResponseIds(Content message) { + return message.parts().stream() + .flatMap(List::stream) + .map(Part::functionResponse) + .flatMap(Optional::stream) + .map(FunctionResponse::id) + .flatMap(Optional::stream) + .collect(toCollection(HashSet::new)); + } + + /** + * Returns the user message that started {@code invocationId}, so a resume that carries no message + * of its own still shows plugins, callbacks and the model the original prompt instead of empty + * content. A function-response message does not qualify: it answers an invocation rather than + * starting one. + * + * @throws IllegalArgumentException if the invocation has no such message, as in Python ADK. + */ + private static Content originalUserMessage(Session session, String invocationId) { + return session.events().stream() + .filter(event -> invocationId.equals(event.invocationId())) + .filter(event -> Objects.equals(event.author(), Role.USER)) + .map(event -> event.content().orElse(null)) + .filter(Objects::nonNull) + .filter(content -> content.parts().stream().flatMap(List::stream).findAny().isPresent()) + .filter( + content -> + content.parts().stream() + .flatMap(List::stream) + .noneMatch(part -> part.functionResponse().isPresent())) + .findFirst() + .orElseThrow( + () -> + new IllegalArgumentException( + "No user message to resume the requested invocation from.")); + } + + /** 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()); @@ -728,6 +1158,10 @@ private InvocationContext newInvocationContextForLive( newInvocationContextBuilder(session) .runConfig(runConfigBuilder.build()) .userContent(Content.fromParts()) + // A live run has no pause to resume, so route to the function call's author as before + // resumability: re-entering the workflow would hand the turn to a LoopAgent or + // ParallelAgent, neither of which implements runLive. + .agent(findAgentToRun(session, this.agent, /* reEnterWorkflow= */ false)) .liveRequestQueue(liveRequestQueue); return builder.build(); @@ -740,7 +1174,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) @@ -832,10 +1265,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)); }); @@ -868,19 +1298,44 @@ private boolean isResumable() { return resumabilityConfig != null && resumabilityConfig.isResumable(); } + /** + * Returns whether this runner's app runs the legacy resumption flow. A separate flow from {@link + * #isResumable()}, so a caller wanting either has to ask for both. + */ + @SuppressWarnings("deprecation") // The shim it reads is deprecated by design. + private boolean isLegacyResumability() { + 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. + return findAgentToRun(session, rootAgent, /* reEnterWorkflow= */ true); + } + + /** + * Returns the agent that should handle the next request based on session history. {@code + * reEnterWorkflow} selects the resume-aware routing: false routes to the function call's author + * itself, as Python always does, for a caller with no resume semantics to serve. + */ + private BaseAgent findAgentToRun(Session session, BaseAgent rootAgent, boolean reEnterWorkflow) { + // Route a function response to its call's author, re-entering via the top-most workflow + // ancestor each mode treats as resume-aware so that workflow can advance past it. 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(); + if (!reEnterWorkflow) { + return author; + } + if (isResumable()) { + return topmostResumableWorkflowAncestor(author); + } + if (isLegacyResumability()) { + return topmostSequentialAncestor(author); + } + return author; } List events = new ArrayList<>(session.events()); @@ -895,6 +1350,16 @@ 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. Only when resumable: endOfAgent shares a field with endInvocation, which + // any callback or tool can set, so skipping it unconditionally would change routing for + // apps that never enabled resumability. + if (isResumable() + && (event.actions().endOfAgent() || event.actions().agentState().isPresent())) { + continue; + } + if (author.equals(rootAgent.name())) { return rootAgent; } @@ -913,11 +1378,31 @@ private BaseAgent findAgentToRun(Session session, BaseAgent rootAgent) { return rootAgent; } + /** + * Returns the top-most ancestor reachable from {@code agent} through resume-aware workflow + * parents ({@link SequentialAgent}, {@link LoopAgent} or {@link ParallelAgent}), 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. Each of the three skips sub-agents already + * marked end-of-agent, so re-entering replays no completed work. Deliberately stricter than + * Python ADK, which routes to the author itself. + */ + private static BaseAgent topmostResumableWorkflowAncestor(BaseAgent agent) { + BaseAgent result = agent; + BaseAgent parent = agent.parentAgent(); + while (parent instanceof SequentialAgent + || parent instanceof LoopAgent + || parent instanceof ParallelAgent) { + result = parent; + parent = parent.parentAgent(); + } + return result; + } + /** * 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). + * parents, or {@code agent} itself otherwise. Used by the deprecated legacy flow, whose only + * resume-aware workflow agent is SequentialAgent; the resumable flow routes to the author itself, + * as Python does. */ private static BaseAgent topmostSequentialAncestor(BaseAgent agent) { BaseAgent result = agent; 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..7aeab233b 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,14 @@ static Event fromApiEvent(Map apiEvent) { Optional.ofNullable(actionsMap.get("requestedToolConfirmations")) .map(SessionJsonConverter::asConcurrentMapOfToolConfirmations) .orElse(new ConcurrentHashMap<>())); + Object agentState = actionsMap.get("agentState"); + if (agentState instanceof Map) { + eventActionsBuilder.agentState((Map) agentState); + } else if (agentState != null) { + // Drop a non-map agentState (a session written by another ADK runtime) rather than failing + // the whole session load on an unchecked cast. + logger.warn("Ignoring 'agentState' of unexpected type {}", agentState.getClass().getName()); + } } 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..ce849a6c7 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; @@ -67,6 +75,46 @@ public void setUp() { activeStreamingTools.put("test-tool", new ActiveStreamingTool(new LiveRequestQueue())); } + // The two flags now select different flows: the shim runs the legacy one, so it must not report + // resumable. Replaces the equivalence the flags had when the shim first started selecting + // resumption. + @Test + @SuppressWarnings("deprecation") // Exercises the deprecated shim. + public void isResumable_andIsLegacyResumability_separateTheTwoModes() { + InvocationContext shimContext = + contextWith(ResumabilityConfig.builder().plainTextContinuationAutoResume(true).build()); + InvocationContext resumableContext = + contextWith(ResumabilityConfig.builder().resumable(true).build()); + InvocationContext explicitlyOffContext = + contextWith(ResumabilityConfig.builder().resumable(false).build()); + InvocationContext neitherContext = contextWith(null); + + assertThat(shimContext.isResumable()).isFalse(); + assertThat(shimContext.isLegacyResumability()).isTrue(); + + assertThat(resumableContext.isResumable()).isTrue(); + assertThat(resumableContext.isLegacyResumability()).isFalse(); + + // An explicit false and an absent config are the same answer, either side of the null guard. + assertThat(explicitlyOffContext.isResumable()).isFalse(); + assertThat(neitherContext.isResumable()).isFalse(); + assertThat(neitherContext.isLegacyResumability()).isFalse(); + } + + private InvocationContext contextWith(ResumabilityConfig resumabilityConfig) { + return InvocationContext.builder() + .sessionService(mockSessionService) + .artifactService(mockArtifactService) + .pluginManager(pluginManager) + .invocationId(testInvocationId) + .agent(mockAgent) + .session(session) + .userContent(userContent) + .runConfig(runConfig) + .resumabilityConfig(resumabilityConfig) + .build(); + } + @Test public void testBuildWithUserContent() { InvocationContext context = @@ -724,4 +772,398 @@ 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 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_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..4d39e2f2a 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,235 @@ 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(); + } + + /** + * A transfer event shaped the way the flow really records one: the {@code transfer_to_agent} + * tool's function response, authored by the transferring agent and carrying the transfer action. + */ + private static Event transferEvent(String author, String targetAgent) { + return Event.builder() + .id("t1") + .invocationId("inv") + .author(author) + .actions(EventActions.builder().transferToAgent(targetAgent).build()) + .content( + Content.builder() + .role("user") + .parts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("transfer-call") + .name("transfer_to_agent") + .response(ImmutableMap.of()) + .build()) + .build()) + .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(); + InvocationContext context = resumableContextWithSeededEvent(root, transferEvent("root", "sub")); + 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(); + InvocationContext context = resumableContextWithSeededEvent(root, transferEvent("root", "sub")); + 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_resumeTransferActionWithoutTransferResponse_continuesRootAgent() { + // Only a transfer_to_agent function response records a transfer. A callback that sets the + // action on some other event must not make the resume run the named agent. + LlmAgent sub = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("sub should not run"))) + .name("sub") + .build(); + TestLlm rootLlm = createTestLlm(createTextLlmResponse("root continues")); + LlmAgent root = createTestAgentBuilder(rootLlm).name("root").subAgents(sub).build(); + Event actionOnPlainText = + Event.builder() + .id("t1") + .invocationId("inv") + .author("root") + .actions(EventActions.builder().transferToAgent("sub").build()) + .content(Content.fromParts(Part.fromText("just text"))) + .build(); + InvocationContext context = resumableContextWithSeededEvent(root, actionOnPlainText); + context.setAgentState("root", ImmutableMap.of(), /* endOfAgent= */ false); + + List events = root.runAsync(context).toList().blockingGet(); + + assertThat(simplifyEvents(events)).contains("root: root continues"); + assertThat(simplifyEvents(events)).doesNotContain("sub: sub should not run"); + } + + @Test + public void runAsync_resumeFromTransferToMissingAgent_throws() { + // The transferred-to agent is gone from the tree. A live transfer fails in the flow; a resumed + // one must fail too, rather than silently re-running the agent that handed off. + TestLlm rootLlm = createTestLlm(createTextLlmResponse("root should not run")); + LlmAgent root = createTestAgentBuilder(rootLlm).name("root").build(); + InvocationContext context = + resumableContextWithSeededEvent(root, transferEvent("root", "removed_agent")); + context.setAgentState("root", ImmutableMap.of(), /* endOfAgent= */ false); + + var resumed = root.runAsync(context).toList(); + assertThrows(IllegalStateException.class, resumed::blockingGet); + assertThat(rootLlm.getRequests()).isEmpty(); + } + + @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(); + InvocationContext context = + resumableContextWithSeededEvent(agentA, transferEvent("agent_a", "agent_b")); + 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..c6d8dff44 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,123 @@ public void runAsync_withEndInvocationInSubAgentCallback_stopsSubAgentButLoopCon assertThat(normalAgentRunCount.get()).isEqualTo(3); assertThat(subAgent2RunCount.get()).isEqualTo(1); } + + // ---- Resumability: durable checkpoint resume. ---- + + @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..97d107aa9 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,195 @@ 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_subAgentEscalates_emitsEndOfAgent() { + // Authored by the sub-agent itself: only a direct sub-agent's escalation addresses this agent. + BaseAgent escalating = + createSubAgent( + "escalating", createEscalateEvent("esc").toBuilder().author("escalating").build()); + 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..77e6fc2e5 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,143 @@ 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(); + } + + // Legacy resumption reconstructs the resume point from history, so it fast-forwards past + // completed sub-agents with no checkpoints in the session. The resumable flow does not: it + // resumes from its own checkpoints, as Python ADK does. + @Test + @SuppressWarnings("deprecation") // The legacy shim is deprecated by design. + public void runAsync_legacyResumption_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(false) + .plainTextContinuationAutoResume(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/apps/ResumabilityConfigTest.java b/core/src/test/java/com/google/adk/apps/ResumabilityConfigTest.java new file mode 100644 index 000000000..2407f5c62 --- /dev/null +++ b/core/src/test/java/com/google/adk/apps/ResumabilityConfigTest.java @@ -0,0 +1,65 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.adk.apps; + +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; + +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests for {@link ResumabilityConfig}. */ +@RunWith(JUnit4.class) +@SuppressWarnings("deprecation") // Exercises the deprecated plainTextContinuationAutoResume shim. +public final class ResumabilityConfigTest { + + @Test + public void build_defaults_selectNoResumption() { + ResumabilityConfig config = ResumabilityConfig.builder().build(); + + assertThat(config.isResumable()).isFalse(); + assertThat(config.isPlainTextContinuationAutoResume()).isFalse(); + } + + @Test + public void build_resumableOnly_succeeds() { + ResumabilityConfig config = ResumabilityConfig.builder().resumable(true).build(); + + assertThat(config.isResumable()).isTrue(); + assertThat(config.isPlainTextContinuationAutoResume()).isFalse(); + } + + @Test + public void build_legacyShimOnly_succeeds() { + ResumabilityConfig config = + ResumabilityConfig.builder().plainTextContinuationAutoResume(true).build(); + + assertThat(config.isResumable()).isFalse(); + assertThat(config.isPlainTextContinuationAutoResume()).isTrue(); + } + + @Test + public void build_bothFlags_throws() { + ResumabilityConfig.Builder builder = + ResumabilityConfig.builder().resumable(true).plainTextContinuationAutoResume(true); + + IllegalArgumentException thrown = assertThrows(IllegalArgumentException.class, builder::build); + + assertThat(thrown).hasMessageThat().contains("mutually exclusive"); + } +} 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..5a31d00dd 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,14 @@ 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(); + } + // 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 +576,38 @@ public void hasPendingLongRunningCall_emptyList_returnsFalse() { assertThat(Functions.hasPendingLongRunningCall(ImmutableList.of())).isFalse(); } + @Test + public void hasPendingLongRunningCall_list_matchingResponse_stillReturnsTrue() { + // Responses are not matched against calls: the call event is inside the two-event window, so + // even its own answer leaves this true. The resumable flow matches responses in StepResume; + // this predicate serves only the legacy flow, which pauses here as it always did. + ImmutableList events = + ImmutableList.of(longRunningCallEvent("call1"), functionResponseEvent("call1")); + 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/RunnerLegacyResumabilityTest.java b/core/src/test/java/com/google/adk/runner/RunnerLegacyResumabilityTest.java new file mode 100644 index 000000000..ad23f92f8 --- /dev/null +++ b/core/src/test/java/com/google/adk/runner/RunnerLegacyResumabilityTest.java @@ -0,0 +1,728 @@ +/* + * 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.runner; + +import static com.google.adk.testing.ResumabilityTestUtils.answerCall; +import static com.google.adk.testing.ResumabilityTestUtils.newSession; +import static com.google.adk.testing.ResumabilityTestUtils.pendingFunctionTool; +import static com.google.adk.testing.ResumabilityTestUtils.runTurn; +import static com.google.adk.testing.ResumabilityTestUtils.shimRunner; +import static com.google.adk.testing.TestUtils.createFunctionCallLlmResponse; +import static com.google.adk.testing.TestUtils.createLlmResponse; +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.ImmutableList.toImmutableList; +import static com.google.common.truth.Truth.assertThat; + +import com.google.adk.agents.LlmAgent; +import com.google.adk.agents.LoopAgent; +import com.google.adk.agents.ParallelAgent; +import com.google.adk.agents.RunConfig; +import com.google.adk.agents.SequentialAgent; +import com.google.adk.events.Event; +import com.google.adk.flows.llmflows.Functions; +import com.google.adk.sessions.Session; +import com.google.adk.telemetry.Tracing; +import com.google.adk.testing.ResumabilityTestUtils.Tools; +import com.google.adk.testing.TestLlm; +import com.google.adk.tools.FunctionTool; +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.FunctionResponse; +import com.google.genai.types.Part; +import io.opentelemetry.api.trace.Tracer; +import io.opentelemetry.sdk.testing.junit4.OpenTelemetryRule; +import io.reactivex.rxjava3.core.Flowable; +import java.util.Objects; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.After; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * Runner tests for the deprecated plain-text continuation shim. + * + *

Includes a copy of every resumability test that predates durable checkpoints, re-run under the + * shim, so the shim keeps behaving exactly as resumability did before it was split in two. + */ +@RunWith(JUnit4.class) +@SuppressWarnings("deprecation") // The class exists to exercise the deprecated shim flag. +public final class RunnerLegacyResumabilityTest { + @Rule public final OpenTelemetryRule openTelemetryRule = OpenTelemetryRule.create(); + + private Tracer originalTracer; + + @Before + public void setUp() { + this.originalTracer = Tracing.getTracer(); + Tracing.setTracerForTesting( + openTelemetryRule.getOpenTelemetry().getTracer("RunnerLegacyResumabilityTest")); + } + + @After + public void tearDown() { + Tracing.setTracerForTesting(originalTracer); + } + + // 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 + 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 = shimRunner(rootAgent); + Session session = newSession(runner); + + // 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 = runTurn(runner, session, "m1"); + for (String expected : new String[] {"r2", "r3", "r4", "r5"}) { + ImmutableList turn = runTurn(runner, session, "m"); + 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 + public void runAsync_plainTextAutoResume_checkpointlessSession_startsNewInvocation() { + LlmAgent agent = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("second answer"))) + .name("agent") + .build(); + Runner runner = shimRunner(agent); + Session session = newSession(runner); + + // 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(); + + ImmutableList secondTurn = runTurn(runner, session, "second turn"); + + // 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(); + } + + // The shim selects the legacy resumption flow but does NOT resume on a plain-text turn: that + // starts a new invocation, as it did before the flag was honored. + @Test + @SuppressWarnings( + "deprecation") // Resumability + the auto-resume shim are intentionally deprecated. + public void runAsync_plainTextContinuation_shimOn_startsNewInvocation() { + TestLlm testLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "draft")), + createTextLlmResponse("continued")); + LlmAgent agent = + createTestAgentBuilder(testLlm).name("agent").tools(pendingFunctionTool()).build(); + Runner runner = shimRunner(agent); + Session session = newSession(runner); + + ImmutableList turn1 = runTurn(runner, session, "draft the note"); + String invocationId = turn1.get(0).invocationId(); + assertThat(testLlm.getRequests()).hasSize(1); // paused after a single model call + + ImmutableList turn2 = runTurn(runner, session, "Proceed"); + + // The shim does not resume on plain text: the continuation opens a NEW invocation and the + // agent plans again, exactly as it did before the shim was honored. + assertThat(testLlm.getRequests()).hasSize(2); + assertThat(turn2).isNotEmpty(); + 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()).isNotEqualTo(invocationId); + assertThat(turn2.stream().noneMatch(event -> invocationId.equals(event.invocationId()))) + .isTrue(); + } + + // 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 = shimRunner(agent); + Session session = newSession(runner); + + ImmutableList turn1 = runTurn(runner, session, "hi"); + String invocationId = turn1.get(0).invocationId(); + assertThat(simplifyEvents(turn1)).contains("agent: first answer"); + + ImmutableList turn2 = runTurn(runner, session, "again"); + + assertThat(simplifyEvents(turn2)).contains("agent: second answer"); + assertThat(turn2.get(0).invocationId()).isNotEqualTo(invocationId); + } + + // Nested topology: a paused long-running call inside an LlmAgent nested in a SequentialAgent is + // treated like the flat case -- the shim does not resume on plain text. + @Test + @SuppressWarnings( + "deprecation") // Resumability + the auto-resume shim are intentionally deprecated. + public void runAsync_plainTextContinuation_inSequentialAgent_shimOn_startsNewInvocation() { + 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(pendingFunctionTool()).build(); + SequentialAgent workflowAgent = + SequentialAgent.builder() + .name("workflow_agent") + .subAgents(ImmutableList.of(childAgent)) + .build(); + Runner runner = shimRunner(workflowAgent); + Session session = newSession(runner); + + ImmutableList turn1 = runTurn(runner, session, "draft the note"); + String invocationId = turn1.get(0).invocationId(); + assertThat(testLlm.getRequests()).hasSize(1); // paused inside the workflow after one model call + + Object unused = runTurn(runner, session, "Proceed"); + + // Nested under a SequentialAgent the shim behaves as it does for a flat agent: the plain-text + // turn opens a new invocation rather than resuming the paused one. + assertThat(testLlm.getRequests()).hasSize(2); + 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()).isNotEqualTo(invocationId); + } + + // Shim-only: a long-running tool that returns a value in the same turn still pauses the flow + // after one model call, as it did before checkpoints. Guards BaseLlmFlow's post-step pause, which + // is now reachable only from the shim -- a value-returning tool is the case that distinguishes it + // from the flow simply ending on a final response. + @Test + @SuppressWarnings( + "deprecation") // Resumability + the auto-resume shim are intentionally deprecated. + public void runAsync_shimOnly_valueReturningLongRunningCall_pausesAfterSingleModelCall() { + TestLlm testLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "echoTool", ImmutableMap.of("message", "hi")), + createTextLlmResponse("should not be reached")); + LlmAgent agent = + createTestAgentBuilder(testLlm) + .name("agent") + .tools( + FunctionTool.create( + Tools.class, + "echoTool", + /* requireConfirmation= */ false, + /* isLongRunning= */ true)) + .build(); + Runner runner = shimRunner(agent); + Session session = newSession(runner); + + ImmutableList events = runTurn(runner, session, "go"); + + // The tool answered in-turn, but the flow still pauses rather than summarizing. + assertThat(testLlm.getRequests()).hasSize(1); + assertThat(simplifyEvents(events)).doesNotContain("agent: should not be reached"); + } + + // Shim-only: the LoopAgent legacy body stops looping once a sub-agent emits a pending + // long-running call, and writes no checkpoint. + @Test + @SuppressWarnings( + "deprecation") // Resumability + the auto-resume shim are intentionally deprecated. + public void runAsync_shimOnlyInLoopAgent_stopsOnPendingCallAndWritesNoCheckpoints() { + TestLlm llm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "draft")), + createTextLlmResponse("second iteration")); + LlmAgent child = + createTestAgentBuilder(llm).name("looped_agent").tools(pendingFunctionTool()).build(); + LoopAgent loopAgent = + LoopAgent.builder().name("loop").maxIterations(3).subAgents(child).build(); + Runner runner = shimRunner(loopAgent); + Session session = newSession(runner); + + var unused = runTurn(runner, session, "start"); + + // The pending call stops the loop after the first iteration, and nothing is checkpointed. + assertThat(llm.getRequests()).hasSize(1); + Session reloaded = + runner + .sessionService() + .getSession("test", "user", session.id(), Optional.empty()) + .blockingGet(); + assertThat( + reloaded.events().stream() + .filter( + event -> + event.actions().endOfAgent() || event.actions().agentState().isPresent()) + .collect(toImmutableList())) + .isEmpty(); + } + + // Shim-only (resumable off, shim on): a long-running pause in the middle sub-agent holds the + // workflow, and the plain-text continuation restarts the sequence -- the pre-checkpoint behavior + // the shim reproduces. The point of the test is that no checkpoint is ever persisted. + @Test + @SuppressWarnings( + "deprecation") // Resumability + the auto-resume shim are intentionally deprecated. + public void runAsync_shimOnlyInSequentialAgent_restartsSequenceAndWritesNoCheckpoints() { + TestLlm llmA = createTestLlm(createTextLlmResponse("A ran"), createTextLlmResponse("A replay")); + TestLlm llmB = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "draft")), + createTextLlmResponse("B done")); + TestLlm llmC = createTestLlm(createTextLlmResponse("C ran")); + LlmAgent agentA = createTestAgentBuilder(llmA).name("agent_a").build(); + LlmAgent agentB = + createTestAgentBuilder(llmB).name("agent_b").tools(pendingFunctionTool()).build(); + LlmAgent agentC = createTestAgentBuilder(llmC).name("agent_c").build(); + SequentialAgent workflowAgent = + SequentialAgent.builder() + .name("workflow_agent") + .subAgents(ImmutableList.of(agentA, agentB, agentC)) + .build(); + Runner runner = shimRunner(workflowAgent); + Session session = newSession(runner); + + ImmutableList turn1 = runTurn(runner, session, "draft the note"); + + // Turn 1 pauses at agent_b: agent_c never runs. + assertThat(llmA.getRequests()).hasSize(1); + assertThat(llmB.getRequests()).hasSize(1); + assertThat(llmC.getRequests()).isEmpty(); + assertThat(simplifyEvents(turn1)).doesNotContain("agent_c: C ran"); + + ImmutableList turn2 = runTurn(runner, session, "Proceed"); + + // No durable state to resume from, so the sequence restarts at agent_a and runs through. + assertThat(llmA.getRequests()).hasSize(2); + assertThat(llmB.getRequests()).hasSize(2); + assertThat(llmC.getRequests()).hasSize(1); + assertThat(simplifyEvents(turn2)).contains("agent_c: C ran"); + + // The shim never writes durable state: nothing was checkpointed into the session. + Session reloaded = + runner + .sessionService() + .getSession("test", "user", session.id(), Optional.empty()) + .blockingGet(); + assertThat( + reloaded.events().stream() + .filter( + event -> + event.actions().endOfAgent() || event.actions().agentState().isPresent()) + .collect(toImmutableList())) + .isEmpty(); + } + + // Shim-only, nested workflow: a grandchild-authored pause behaves like the flat case -- the + // plain-text continuation restarts the sequence rather than resuming into the paused sub-agent. + @Test + @SuppressWarnings( + "deprecation") // Resumability + the auto-resume shim are intentionally deprecated. + public void runAsync_shimOnlyInNestedSequentialAgent_restartsSequence() { + TestLlm llmA = createTestLlm(createTextLlmResponse("A ran"), createTextLlmResponse("A replay")); + TestLlm llmB = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "draft")), + createTextLlmResponse("B done")); + LlmAgent agentA = createTestAgentBuilder(llmA).name("agent_a").build(); + LlmAgent agentB = + createTestAgentBuilder(llmB).name("agent_b").tools(pendingFunctionTool()).build(); + SequentialAgent innerWorkflow = + SequentialAgent.builder() + .name("inner_workflow") + .subAgents(ImmutableList.of(agentB)) + .build(); + SequentialAgent workflowAgent = + SequentialAgent.builder() + .name("workflow_agent") + .subAgents(ImmutableList.of(agentA, innerWorkflow)) + .build(); + Runner runner = shimRunner(workflowAgent); + Session session = newSession(runner); + + var unused = runTurn(runner, session, "draft the note"); + ImmutableList turn2 = runTurn(runner, session, "Proceed"); + + assertThat(llmA.getRequests()).hasSize(2); // the sequence restarts at agent_a + assertThat(llmB.getRequests()).hasSize(2); + assertThat(simplifyEvents(turn2)).contains("agent_b: B done"); + } + + // With the deprecated shim, a plain-text continuation starts a new invocation and a non-null + // stateDelta is still merged into the session. + @Test + public void runAsync_plainTextWithShim_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(pendingFunctionTool()).build(); + Runner runner = shimRunner(agent); + Session session = newSession(runner); + + ImmutableList pausedTurn = runTurn(runner, session, "start"); + 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); + ImmutableList resumed = + ImmutableList.copyOf( + runner + .runAsync( + "user", + session.id(), + Content.fromParts(Part.fromText("Proceed")), + RunConfig.builder().build(), + stateDelta) + .toList() + .blockingGet()); + + // The shim starts a new invocation for the plain-text turn; the stateDelta still lands. + assertThat(testLlm.getRequests()).hasSize(2); + assertThat(resumed).isNotEmpty(); + 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()).isNotEqualTo(pausedInvocationId); + assertThat(continuation.actions().stateDelta()).containsAtLeastEntriesIn(stateDelta); + assertThat(finalSession.state()).containsAtLeastEntriesIn(stateDelta); + } + + // ===== CL1-parity: every CL1 resumable(true) test, re-run under the text-only shim ===== + + @Test + public void + runAsync_withToolConfirmation_inSequentialAgent_runsLaterSubAgentsAfterResume_legacyShim() { + LlmAgent agentA = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("agent A done"))) + .name("a_agent") + .build(); + // With resumability on, B pauses right after requesting confirmation (no extra model call), so + // a + // single follow-up response covers the resume. + TestLlm bTestLlm = + createTestLlm( + createFunctionCallLlmResponse( + "tool_call_id", "echoTool", ImmutableMap.of("message", "hello")), + createTextLlmResponse("Response after user confirmed.")); + LlmAgent agentB = + createTestAgentBuilder(bTestLlm) + .name("b_agent") + .tools(FunctionTool.create(Tools.class, "echoTool", /* requireConfirmation= */ 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 = shimRunner(workflowAgent); + Session session = newSession(runner); + + ImmutableList eventsBeforeConfirmation = runTurn(runner, session, "from user"); + + // Turn 1: A runs, B pauses for confirmation, and C must not run yet. + assertThat(simplifyEvents(eventsBeforeConfirmation)).contains("a_agent: agent A done"); + assertThat(simplifyEvents(eventsBeforeConfirmation)).doesNotContain("c_agent: agent C done"); + + FunctionCall askUserConfirmationFunctionCall = + Iterables.getOnlyElement( + eventsBeforeConfirmation.stream() + .map(Functions::getAskUserConfirmationFunctionCalls) + .filter(functionCalls -> !functionCalls.isEmpty()) + .findFirst() + .get()); + ImmutableList eventsAfterConfirmation = + ImmutableList.copyOf( + runner + .runAsync( + "user", + session.id(), + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id(askUserConfirmationFunctionCall.id().get()) + .name(askUserConfirmationFunctionCall.name().get()) + .response(ImmutableMap.of("confirmed", true))) + .build())) + .toList() + .blockingGet()); + + // Turn 2: B resumes and executes the tool, then C runs. A is not re-run. + assertThat(simplifyEvents(eventsAfterConfirmation)) + .containsExactly( + "b_agent: FunctionResponse(name=echoTool, response={message=hello})", + "b_agent: Response after user confirmed.", + "c_agent: agent C done") + .inOrder(); + } + + @Test + public void + runAsync_withLongRunningCall_inSequentialAgent_runsLaterSubAgentsAfterResume_legacyShim() { + LlmAgent agentA = + 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. + TestLlm bTestLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "echoTool", ImmutableMap.of("message", "hello")), + createTextLlmResponse("agent B resumed")); + LlmAgent agentB = + createTestAgentBuilder(bTestLlm) + .name("b_agent") + .tools( + FunctionTool.create( + Tools.class, + "echoTool", + /* 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 = shimRunner(workflowAgent); + Session session = newSession(runner); + + ImmutableList eventsBeforeResume = runTurn(runner, session, "from user"); + + // 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"); + + ImmutableList eventsAfterResume = + answerCall(runner, session, "lro_call_id", "echoTool", ImmutableMap.of("message", "hello")); + + // 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(); + } + + @Test + public void runAsync_withLongRunningCall_resumable_pausesAfterSingleModelCall_legacyShim() { + 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(); + Runner runner = shimRunner(agent); + Session session = newSession(runner); + + ImmutableList events = runTurn(runner, session, "from user"); + + // 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"); + } + + @Test + public void + runAsync_loopAgentWithLongRunningSubAgent_resumable_stopsAfterFirstIteration_legacyShim() { + AtomicInteger calls = new AtomicInteger(); + TestLlm loopLlm = + createTestLlm( + () -> + calls.incrementAndGet() <= 5 + ? Flowable.just( + createFunctionCallLlmResponse( + "lro_call_id", "echoTool", ImmutableMap.of("message", "hello"))) + : Flowable.just(createTextLlmResponse("stop"))); + LlmAgent inner = + createTestAgentBuilder(loopLlm) + .name("inner") + .tools( + FunctionTool.create( + Tools.class, + "echoTool", + /* requireConfirmation= */ false, + /* isLongRunning= */ true)) + .build(); + LoopAgent loop = + LoopAgent.builder() + .name("loop") + .subAgents(ImmutableList.of(inner)) + .maxIterations(3) + .build(); + Runner runner = shimRunner(loop); + Session session = newSession(runner); + + ImmutableList unused = runTurn(runner, session, "from user"); + + // Paused after the first iteration: one model call, not maxIterations. + assertThat(loopLlm.getRequests()).hasSize(1); + } + + @Test + public void + runAsync_parallelAgentWithLongRunningBranch_resumable_otherBranchCompletes_legacyShim() { + TestLlm longRunningLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "echoTool", ImmutableMap.of("message", "hello")), + createTextLlmResponse("unexpected")); + LlmAgent longRunningBranch = + createTestAgentBuilder(longRunningLlm) + .name("long_running_branch") + .tools( + FunctionTool.create( + Tools.class, + "echoTool", + /* requireConfirmation= */ false, + /* isLongRunning= */ true)) + .build(); + LlmAgent plainBranch = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("plain branch done"))) + .name("plain_branch") + .build(); + ParallelAgent parallel = + ParallelAgent.builder() + .name("parallel") + .subAgents(ImmutableList.of(longRunningBranch, plainBranch)) + .build(); + Runner runner = shimRunner(parallel); + Session session = newSession(runner); + + ImmutableList events = runTurn(runner, session, "from user"); + + // The long-running branch paused after one model call; the other branch still completed. + assertThat(longRunningLlm.getRequests()).hasSize(1); + assertThat(simplifyEvents(events)).contains("plain_branch: plain branch done"); + } +} diff --git a/core/src/test/java/com/google/adk/runner/RunnerResumabilityTest.java b/core/src/test/java/com/google/adk/runner/RunnerResumabilityTest.java new file mode 100644 index 000000000..2e9446cc0 --- /dev/null +++ b/core/src/test/java/com/google/adk/runner/RunnerResumabilityTest.java @@ -0,0 +1,1514 @@ +/* + * 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.runner; + +import static com.google.adk.testing.ResumabilityTestUtils.answerCall; +import static com.google.adk.testing.ResumabilityTestUtils.assertEndOfAgent; +import static com.google.adk.testing.ResumabilityTestUtils.assertNoEndOfAgent; +import static com.google.adk.testing.ResumabilityTestUtils.confirmingEchoFunctionTool; +import static com.google.adk.testing.ResumabilityTestUtils.longRunningEchoFunctionTool; +import static com.google.adk.testing.ResumabilityTestUtils.newSession; +import static com.google.adk.testing.ResumabilityTestUtils.pauseThenSay; +import static com.google.adk.testing.ResumabilityTestUtils.pausingAgent; +import static com.google.adk.testing.ResumabilityTestUtils.pendingFunctionTool; +import static com.google.adk.testing.ResumabilityTestUtils.resumableRunner; +import static com.google.adk.testing.ResumabilityTestUtils.resume; +import static com.google.adk.testing.ResumabilityTestUtils.runTurn; +import static com.google.adk.testing.ResumabilityTestUtils.textAgent; +import static com.google.adk.testing.TestUtils.createFunctionCallLlmResponse; +import static com.google.adk.testing.TestUtils.createLlmResponse; +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.adk.testing.TestUtils.simplifyResumableEvents; +import static com.google.common.collect.ImmutableList.toImmutableList; +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; +import static org.mockito.ArgumentMatchers.any; +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.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.google.adk.agents.InvocationContext; +import com.google.adk.agents.LlmAgent; +import com.google.adk.agents.LoopAgent; +import com.google.adk.agents.ParallelAgent; +import com.google.adk.agents.RunConfig; +import com.google.adk.agents.SequentialAgent; +import com.google.adk.apps.App; +import com.google.adk.events.Event; +import com.google.adk.events.EventActions; +import com.google.adk.flows.llmflows.Functions; +import com.google.adk.plugins.BasePlugin; +import com.google.adk.sessions.Session; +import com.google.adk.telemetry.Tracing; +import com.google.adk.testing.ResumabilityTestUtils.Tools; +import com.google.adk.testing.TestBaseAgent; +import com.google.adk.testing.TestLlm; +import com.google.adk.tools.FunctionTool; +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.FunctionResponse; +import com.google.genai.types.Part; +import io.opentelemetry.api.trace.Tracer; +import io.opentelemetry.sdk.testing.junit4.OpenTelemetryRule; +import io.reactivex.rxjava3.core.Flowable; +import io.reactivex.rxjava3.core.Maybe; +import java.util.List; +import java.util.Objects; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicInteger; +import org.jspecify.annotations.Nullable; +import org.junit.After; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Runner tests for durable resumability: {@code ResumabilityConfig.resumable(true)}. */ +@RunWith(JUnit4.class) +@SuppressWarnings("deprecation") // The class exists to exercise the deprecated resumability flags. +public final class RunnerResumabilityTest { + @Rule public final OpenTelemetryRule openTelemetryRule = OpenTelemetryRule.create(); + + private Tracer originalTracer; + + private BasePlugin mockPlugin(String name) { + // Need CALLS_REAL_METHODS to avoid NPE. The default implementation is only returning + // Maybe.empty() + BasePlugin plugin = mock(BasePlugin.class, CALLS_REAL_METHODS); + when(plugin.getName()).thenReturn(name); + return plugin; + } + + @Before + public void setUp() { + this.originalTracer = Tracing.getTracer(); + Tracing.setTracerForTesting( + openTelemetryRule.getOpenTelemetry().getTracer("RunnerResumabilityTest")); + } + + @After + public void tearDown() { + Tracing.setTracerForTesting(originalTracer); + } + + // OSS HITL: after an adk_request_confirmation resumes sub-agent B in a SequentialAgent(A, B, C), + // the workflow must advance to C without re-running the already completed A. + @Test + public void runAsync_withToolConfirmation_inSequentialAgent_runsLaterSubAgentsAfterResume() { + LlmAgent agentA = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("agent A done"))) + .name("a_agent") + .build(); + // With resumability on, B pauses right after requesting confirmation (no extra model call), so + // a + // single follow-up response covers the resume. + TestLlm bTestLlm = + createTestLlm( + createFunctionCallLlmResponse( + "tool_call_id", "echoTool", ImmutableMap.of("message", "hello")), + createTextLlmResponse("Response after user confirmed.")); + LlmAgent agentB = + createTestAgentBuilder(bTestLlm) + .name("b_agent") + .tools(confirmingEchoFunctionTool()) + .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 = resumableRunner(workflowAgent); + Session session = newSession(runner); + + ImmutableList eventsBeforeConfirmation = runTurn(runner, session, "from user"); + + // Turn 1: A runs, B pauses for confirmation, and C must not run yet. + assertThat(simplifyEvents(eventsBeforeConfirmation)).contains("a_agent: agent A done"); + assertThat(simplifyEvents(eventsBeforeConfirmation)).doesNotContain("c_agent: agent C done"); + + FunctionCall askUserConfirmationFunctionCall = + Iterables.getOnlyElement( + eventsBeforeConfirmation.stream() + .map(Functions::getAskUserConfirmationFunctionCalls) + .filter(functionCalls -> !functionCalls.isEmpty()) + .findFirst() + .get()); + ImmutableList eventsAfterConfirmation = + ImmutableList.copyOf( + runner + .runAsync( + "user", + session.id(), + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id(askUserConfirmationFunctionCall.id().get()) + .name(askUserConfirmationFunctionCall.name().get()) + .response(ImmutableMap.of("confirmed", true))) + .build())) + .toList() + .blockingGet()); + + // 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.", + "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(); + } + + // Long-running-call HITL: a pending long-running function call (not the confirmation flow) pauses + // SequentialAgent(A, B, C) after B; on resume B continues and C runs, without re-running A. + @Test + public void runAsync_withLongRunningCall_inSequentialAgent_runsLaterSubAgentsAfterResume() { + LlmAgent agentA = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("agent A done"))) + .name("a_agent") + .build(); + // 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", "pendingTool", ImmutableMap.of("message", "hello")), + createTextLlmResponse("agent B resumed")); + LlmAgent agentB = + createTestAgentBuilder(bTestLlm).name("b_agent").tools(pendingFunctionTool()).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 = resumableRunner(workflowAgent); + Session session = newSession(runner); + + ImmutableList eventsBeforeResume = runTurn(runner, session, "from user"); + + // 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"); + + ImmutableList eventsAfterResume = + answerCall( + runner, session, "lro_call_id", "pendingTool", ImmutableMap.of("message", "hello")); + + // 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 + 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(pendingFunctionTool()).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 = resumableRunner(workflowAgent); + Session session = newSession(runner); + + ImmutableList eventsBeforeResume = runTurn(runner, session, "from user"); + + // 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"); + + ImmutableList eventsAfterResume = + answerCall( + runner, session, "lro_call_id", "pendingTool", ImmutableMap.of("message", "hello")); + + // 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 + 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 = resumableRunner(rootAgent); + Session session = newSession(runner); + + ImmutableList turn1 = runTurn(runner, session, "hi"); + + // 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(); + + ImmutableList turn2 = runTurn(runner, session, "again"); + + // 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 + 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(pendingFunctionTool()).build(); + LlmAgent rootAgent = + createTestAgentBuilder(testLlm) + .name("root_agent") + .subAgents(ImmutableList.of(subAgent1)) + .build(); + Runner runner = resumableRunner(rootAgent); + Session session = newSession(runner); + + ImmutableList pausedTurn = runTurn(runner, session, "hi"); + 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"); + + ImmutableList resumed = + ImmutableList.copyOf( + 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(); + } + + // 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 + 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(pendingFunctionTool()).build(); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + + runTurn(runner, session, "start"); + // Turn 1: both long-running calls are issued and the invocation pauses; the model is called + // once. + assertThat(testLlm.getRequests()).hasSize(1); + + ImmutableList afterFirstAnswer = + answerCall(runner, session, "call_a", "pendingTool", ImmutableMap.of("message", "a")); + // One call answered is not enough: nothing runs and the model is not re-invoked. + assertThat(afterFirstAnswer).isEmpty(); + assertThat(testLlm.getRequests()).hasSize(1); + + ImmutableList afterSecondAnswer = + answerCall(runner, session, "call_b", "pendingTool", ImmutableMap.of("message", "b")); + // 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 + 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(longRunningEchoFunctionTool()).build(); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + + ImmutableList events = runTurn(runner, session, "from user"); + + // The value-returning call was summarized in the same turn: the model was re-invoked (two + // calls) and the summary surfaced, with no pause. No end-of-agent checkpoint is emitted, so + // the invocation stays resumable. + assertThat(testLlm.getRequests()).hasSize(2); + 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 + 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(pendingFunctionTool()).build(); + Runner runner = resumableRunner(agent, resumePlugin); + Session session = newSession(runner); + + runTurn(runner, session, "start"); + + ImmutableList resumed = + answerCall( + runner, session, "lro_call_id", "pendingTool", ImmutableMap.of("message", "hello")); + + 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()); + } + + // An unanswered long-running call keeps the invocation paused however it is resumed. A plain-text + // message continues the agent instead of throwing "No matching function call", and a resume with + // no message at all is equally inert; in neither case is the model re-invoked. + @Test + public void resume_pausedCallUnanswered_staysPausedForPlainTextAndNoMessage() { + TestLlm testLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hello")), + createTextLlmResponse("should not be reached")); + LlmAgent agent = + createTestAgentBuilder(testLlm).name("agent").tools(pendingFunctionTool()).build(); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + + ImmutableList firstTurn = runTurn(runner, session, "start"); + String invocationId = firstTurn.get(0).invocationId(); + + // A plain-text message with an explicit invocation id: before the fix this threw + // IllegalArgumentException; now it returns without throwing. + ImmutableList resumedWithText = + resume( + runner, + session, + invocationId, + Content.fromParts(Part.fromText("please continue")), + /* stateDelta= */ null); + assertThat(resumedWithText).isEmpty(); + + ImmutableList resumedWithNoMessage = + resume(runner, session, invocationId, /* newMessage= */ null, /* stateDelta= */ null); + assertThat(resumedWithNoMessage).isEmpty(); + // The call is still unanswered, so the model was never re-invoked past the pausing turn. + assertThat(testLlm.getRequests()).hasSize(1); + } + + // A nested Sequential whose inner sub-agent pauses silently must leave the OUTER checkpoint on + // the unfinished inner workflow, not advance it -- otherwise the resume re-runs completed work. + @Test + public void runAsync_resume_nestedSequentialSilentPause_checkpointStaysOnUnfinishedInner() { + LlmAgent a1 = pausingAgent("a1_agent", pauseThenSay("lro_call_id", "a1 resumed")); + LlmAgent a2 = textAgent("a2_agent", "a2 done"); + SequentialAgent inner = + SequentialAgent.builder().name("inner_agent").subAgents(ImmutableList.of(a1, a2)).build(); + LlmAgent b = textAgent("b_agent", "b done"); + SequentialAgent outer = + SequentialAgent.builder().name("outer_agent").subAgents(ImmutableList.of(inner, b)).build(); + Runner runner = resumableRunner(outer); + Session session = newSession(runner); + + ImmutableList turn1 = runTurn(runner, session, "start"); + // a1 pauses, so neither a2 nor b runs, and no workflow closes. + assertThat(simplifyEvents(turn1)).doesNotContain("a2_agent: a2 done"); + assertThat(simplifyEvents(turn1)).doesNotContain("b_agent: b done"); + assertNoEndOfAgent(turn1, "outer_agent"); + // The outer checkpoint still names the unfinished inner workflow. + assertThat(simplifyResumableEvents(turn1)) + .contains("outer_agent: agent_state={current_sub_agent=inner_agent}"); + + ImmutableList resumed = + answerCall(runner, session, "lro_call_id", "pendingTool", ImmutableMap.of("message", "hi")); + + // The resume finishes the inner workflow and only then advances the outer one; nothing that + // already completed is re-run. + assertThat(simplifyEvents(resumed)) + .containsAtLeast("a1_agent: a1 resumed", "a2_agent: a2 done", "b_agent: b done") + .inOrder(); + assertEndOfAgent(resumed, "outer_agent"); + } + + // A leaf paused under a ParallelAgent resumes, and the enclosing SequentialAgent advances past + // the completed parallel block -- the nesting must not defeat re-entering the workflow. + @Test + public void runAsync_resume_pausedUnderParallelAgent_advancesEnclosingSequential() { + LlmAgent leaf = + createTestAgentBuilder( + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hi")), + createTextLlmResponse("leaf resumed"))) + .name("leaf_agent") + .tools(pendingFunctionTool()) + .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 = resumableRunner(root); + Session session = newSession(runner); + + runTurn(runner, session, "start"); + + ImmutableList resumed = + answerCall(runner, session, "lro_call_id", "pendingTool", ImmutableMap.of("message", "hi")); + + // The leaf itself resumes and emits its post-tool response. + assertThat(simplifyEvents(resumed)).contains("leaf_agent: leaf resumed"); + // The enclosing SequentialAgent then advances to the sub-agent after the parallel block. + assertThat(simplifyEvents(resumed)).contains("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 + 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(pendingFunctionTool()) + .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 = resumableRunner(root); + Session session = newSession(runner); + + runTurn(runner, session, "start"); + + // Answer only c1; c2 remains unanswered. + ImmutableList resumed = + answerCall(runner, session, "c1", "pendingTool", ImmutableMap.of("message", "a")); + + // 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(); + } + + // 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 + 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(pendingFunctionTool()).build(); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + + ImmutableList turn1 = runTurn(runner, session, "draft the note"); + String invocationId = turn1.get(0).invocationId(); + + ImmutableList turn2 = runTurn(runner, session, "Proceed"); + + assertThat(testLlm.getRequests()).hasSize(2); // new invocation re-invoked the model + assertThat(turn2.get(0).invocationId()).isNotEqualTo(invocationId); + } + + // 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 = newSession(runner); + + ImmutableList events = runTurn(runner, session, "from user"); + + assertThat(events.stream().anyMatch(event -> event.actions().endOfAgent())).isFalse(); + } + + // A resumable LlmAgent that completes normally emits a trailing end-of-agent checkpoint, and + // resuming past it is a no-op: the active agent already ended, so nothing re-runs. + @Test + public void resume_completedInvocation_emitsEndOfAgentThenIsNoOp() { + LlmAgent agent = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("all done"))) + .name("agent") + .build(); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + + ImmutableList firstTurn = runTurn(runner, session, "from user"); + + // The completed agent closes with an end-of-agent checkpoint, so a later run can tell the + // invocation finished. + Event last = Iterables.getLast(firstTurn); + assertThat(last.author()).isEqualTo("agent"); + assertThat(last.actions().endOfAgent()).isTrue(); + + ImmutableList resumed = + resume( + runner, + session, + firstTurn.get(0).invocationId(), + /* newMessage= */ null, + /* stateDelta= */ null); + + assertThat(resumed).isEmpty(); + } + + // A function-response resume continues the SAME invocation that issued the matching call rather + // than minting a new one, and merges a non-null stateDelta into the session as the + // new-invocation path does. + @Test + public void resume_withFunctionResponseAndStateDelta_resumesSameInvocationAndMergesState() { + TestLlm testLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hello")), + createTextLlmResponse("resumed answer")); + LlmAgent agent = + createTestAgentBuilder(testLlm).name("agent").tools(pendingFunctionTool()).build(); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + + ImmutableList pausedTurn = runTurn(runner, session, "from user"); + String pausedInvocationId = pausedTurn.get(0).invocationId(); + assertThat(simplifyEvents(pausedTurn)).doesNotContain("agent: resumed answer"); + + ImmutableMap stateDelta = ImmutableMap.of("key1", "value1", "key2", 42); + ImmutableList resumed = + ImmutableList.copyOf( + 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()); + + // 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(); + 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); + } + + // Python parity (REPLAY_CALLS): a branch whose last event carries a call the previous run never + // executed runs that call on resume, instead of asking the model again. + @Test + public void resume_unexecutedFunctionCall_replaysCallInsteadOfCallingModel() { + TestLlm testLlm = createTestLlm(createTextLlmResponse("should not be reached")); + LlmAgent agent = + createTestAgentBuilder(testLlm) + .name("agent") + .tools(FunctionTool.create(Tools.class, "echoTool")) + .build(); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + String invocationId = "inv-replay"; + var unusedUser = + runner + .sessionService() + .appendEvent( + session, + Event.builder() + .id("u1") + .invocationId(invocationId) + .author("user") + .content(Content.fromParts(Part.fromText("go"))) + .build()) + .blockingGet(); + // The agent's call was persisted but never executed: no function response follows it. + var unusedCall = + runner + .sessionService() + .appendEvent( + session, + Event.builder() + .id("c1") + .invocationId(invocationId) + .author("agent") + .content( + Content.fromParts( + Part.builder() + .functionCall( + FunctionCall.builder() + .id("call_1") + .name("echoTool") + .args(ImmutableMap.of("message", "hi")) + .build()) + .build())) + .build()) + .blockingGet(); + + ImmutableList resumed = + resume(runner, session, invocationId, /* newMessage= */ null, /* stateDelta= */ null); + + // The persisted call was executed rather than re-requested: the scripted model only ever + // returns text, so an echoTool response can only come from replaying the stored call. + assertThat( + resumed.stream() + .flatMap(event -> event.functionResponses().stream()) + .map(response -> response.name().orElse("")) + .collect(toImmutableList())) + .contains("echoTool"); + // Replay happens before the model is consulted, so the first request already carries the + // tool's response rather than asking the model to produce the call again. + assertThat(testLlm.getRequests().get(0).contents().toString()).contains("echoTool"); + } + + // Python parity: resuming with a stateDelta but no message has no user event to carry the delta, + // so it is persisted as a content-less event instead of being kept in memory only. + @Test + public void resume_withStateDeltaAndNoMessage_persistsStateDeltaEvent() { + TestLlm testLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "draft")), + createTextLlmResponse("done")); + LlmAgent agent = + createTestAgentBuilder(testLlm).name("agent").tools(pendingFunctionTool()).build(); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + ImmutableList turn1 = runTurn(runner, session, "start"); + String invocationId = turn1.get(0).invocationId(); + + ImmutableMap stateDelta = ImmutableMap.of("key1", "value1"); + var unused = resume(runner, session, invocationId, /* newMessage= */ null, stateDelta); + + Session reloaded = + runner + .sessionService() + .getSession("test", "user", session.id(), Optional.empty()) + .blockingGet(); + assertThat(reloaded.state()).containsAtLeastEntriesIn(stateDelta); + assertThat( + reloaded.events().stream() + .anyMatch( + event -> + event.content().isEmpty() + && event.actions().stateDelta().containsKey("key1"))) + .isTrue(); + } + + // 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 + 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(pendingFunctionTool()).build(); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + + // Invocation 1 pauses on the long-running call. + ImmutableList inv1 = runTurn(runner, session, "q1"); + String inv1Id = inv1.get(0).invocationId(); + // Invocation 2 finishes; invocation 3 pauses again. + runTurn(runner, session, "q2"); + runTurn(runner, session, "q3"); + + // Resume invocation 1 (the oldest, not the latest) via its function response. + ImmutableList resumed = + ImmutableList.copyOf( + runner + .runAsync( + "user", + session.id(), + inv1Id, + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("call-1") + .name("pendingTool") + .response(ImmutableMap.of("message", "hi"))) + .build()), + RunConfig.builder().build(), + /* stateDelta= */ null) + .toList() + .blockingGet()); + + assertThat(simplifyEvents(resumed)).contains("agent: llm response after resuming invocation 1"); + assertThat(resumed.stream().allMatch(event -> event.invocationId().equals(inv1Id))).isTrue(); + } + + // InMemoryRunnerTest parity: resume by invocationId rehydrates the agent's checkpoint state from + // history so the running agent observes it. + @Test + public void resume_restoresAgentStateFromHistory() { + TestBaseAgent agent = + new TestBaseAgent( + "test_agent", + "desc", + () -> Flowable.empty(), + /* subAgents= */ null, + /* beforeAgentCallbacks= */ null, + /* afterAgentCallbacks= */ null); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + 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(); + + Object unused = + runner + .runAsync( + "user", + session.id(), + "test-inv", + /* newMessage= */ null, + RunConfig.builder().build(), + /* stateDelta= */ null) + .toList() + .blockingGet(); + + assertThat(agent.getLastInvocationContext().agentStates()) + .containsEntry("test_agent", ImmutableMap.of("saved", "state")); + } + + // InMemoryRunnerTest parity: resume by invocationId with a new user message appends that content + // under the resumed invocation. + @Test + public void resume_withNewMessage_appendsUserContentUnderResumedInvocation() { + TestBaseAgent agent = + new TestBaseAgent("test_agent", "desc", () -> Flowable.empty(), null, null, null); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + Object unusedUser = + runner + .sessionService() + .appendEvent( + session, + Event.builder() + .id("u1") + .invocationId("test-inv") + .author("user") + .content(Content.fromParts(Part.fromText("hi"))) + .build()) + .blockingGet(); + + Object unused = + runner + .runAsync( + "user", + session.id(), + "test-inv", + Content.fromParts(Part.fromText("New message")), + RunConfig.builder().build(), + /* stateDelta= */ null) + .toList() + .blockingGet(); + + 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"); + } + + // RunnerTest parity (disabled counterpart): resuming a non-resumable app throws. + @Test + public void resume_notResumable_throwsException() { + LlmAgent agent = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("x"))).name("agent").build(); + Runner runner = + Runner.builder().app(App.builder().name("test").rootAgent(agent).build()).build(); + Session session = newSession(runner); + String sessionId = session.id(); + + 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: + // With NO invocation id, the runner would otherwise resolve the message to a new invocation; + // it must reject the orphan function response instead of feeding it to the model. Paired with + // resume_orphanFunctionResponseWithProvidedInvocationId_throwsIllegalArgument, which covers the + // other entry point -- the two reach the check by different routes. + @Test + public void resume_functionResponseWithNoMatchingCall_throwsIllegalArgument() { + LlmAgent agent = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("x"))).name("agent").build(); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + + 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); + } + + // A message mixing text with a function response is ambiguous -- the response resumes an + // invocation while the text would start a new one -- so the runner rejects it, as Python does. + @Test + public void resume_functionResponseMixedWithText_throwsIllegalArgument() { + LlmAgent agent = + createTestAgentBuilder( + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hi")), + createTextLlmResponse("done"))) + .name("agent") + .tools(pendingFunctionTool()) + .build(); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + runTurn(runner, session, "start"); + + Content mixed = + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("lro_call_id") + .name("pendingTool") + .response(ImmutableMap.of("message", "hi"))) + .build(), + Part.fromText("and also this")); + + runner + .runAsync( + "user", + session.id(), + /* invocationId= */ null, + mixed, + RunConfig.builder().build(), + /* stateDelta= */ null) + .test() + .assertError(IllegalArgumentException.class); + } + + // Function responses answering calls from two different invocations cannot resume either one, so + // the runner rejects the message instead of silently resuming just the newest invocation. + @Test + public void resume_functionResponsesSpanningTwoInvocations_throwsIllegalArgument() { + LlmAgent agent = + createTestAgentBuilder( + createTestLlm( + createFunctionCallLlmResponse( + "call_a", "pendingTool", ImmutableMap.of("message", "a")), + createFunctionCallLlmResponse( + "call_b", "pendingTool", ImmutableMap.of("message", "b")))) + .name("agent") + .tools(pendingFunctionTool()) + .build(); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + + ImmutableList turn1 = runTurn(runner, session, "start"); + // Plain text starts a second invocation, which pauses on a long-running call of its own. + ImmutableList turn2 = runTurn(runner, session, "another"); + assertThat(turn2.get(0).invocationId()).isNotEqualTo(turn1.get(0).invocationId()); + + Content answersBoth = + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("call_a") + .name("pendingTool") + .response(ImmutableMap.of("message", "a"))) + .build(), + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("call_b") + .name("pendingTool") + .response(ImmutableMap.of("message", "b"))) + .build()); + + runner + .runAsync( + "user", + session.id(), + /* invocationId= */ null, + answersBoth, + RunConfig.builder().build(), + /* stateDelta= */ null) + .test() + .assertError(IllegalArgumentException.class); + } + + // Resuming by invocation id alone carries no message, so the runner recovers the invocation's + // original user message: callbacks and the model must not see empty user content. + @Test + public void resume_withNoNewMessage_restoresOriginalUserContent() { + LlmAgent agent = + createTestAgentBuilder( + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hi")), + createTextLlmResponse("done"))) + .name("agent") + .tools(pendingFunctionTool()) + .build(); + UserContentCapturingPlugin plugin = new UserContentCapturingPlugin(); + Runner runner = resumableRunner(agent, plugin); + Session session = newSession(runner); + + ImmutableList pausedTurn = runTurn(runner, session, "the original ask"); + String pausedInvocationId = pausedTurn.get(0).invocationId(); + + resume(runner, session, pausedInvocationId, /* newMessage= */ null, /* stateDelta= */ null); + + assertThat(plugin.lastUserText).isEqualTo("the original ask"); + } + + // A leaf paused on two long-running calls under a SequentialAgent, answered once, stays paused: + // the leaf pauses without emitting anything, and the sequence must not read that as completion. + @Test + public void runAsync_resume_pausedUnderSequential_partiallyAnswered_doesNotAdvance() { + 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(pendingFunctionTool()) + .build(); + LlmAgent next = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("next ran"))) + .name("next_agent") + .build(); + SequentialAgent root = + SequentialAgent.builder().name("seq_agent").subAgents(ImmutableList.of(leaf, next)).build(); + Runner runner = resumableRunner(root); + Session session = newSession(runner); + + runTurn(runner, session, "start"); + + // Answer only c1; c2 remains unanswered. + ImmutableList resumed = + answerCall(runner, session, "c1", "pendingTool", ImmutableMap.of("message", "a")); + + // The sequence holds: neither the leaf's summary nor the following sub-agent runs. + assertThat(simplifyEvents(resumed)).doesNotContain("leaf_agent: leaf summary"); + assertThat(simplifyEvents(resumed)).doesNotContain("next_agent: next ran"); + } + + // An LlmAgent whose transferred-to ParallelAgent holds the paused leaf: the HITL answer comes + // back on the parallel branch, which the LlmAgent's own branch cannot see. It must still resume. + @Test + public void runAsync_resume_hitlAnswerOnSubBranch_resumesInsteadOfStalling() { + LlmAgent leaf = + createTestAgentBuilder( + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hi")), + createTextLlmResponse("leaf resumed"))) + .name("leaf_agent") + .tools(pendingFunctionTool()) + .build(); + ParallelAgent parallel = + ParallelAgent.builder().name("parallel_agent").subAgents(ImmutableList.of(leaf)).build(); + LlmAgent root = + createTestAgentBuilder( + createTestLlm( + createFunctionCallLlmResponse( + "transfer_call", + "transfer_to_agent", + ImmutableMap.of("agent_name", "parallel_agent")))) + .name("root_agent") + .subAgents(parallel) + .build(); + Runner runner = resumableRunner(root); + Session session = newSession(runner); + + ImmutableList firstTurn = runTurn(runner, session, "start"); + assertThat(simplifyEvents(firstTurn)).isNotEmpty(); + + ImmutableList resumed = + answerCall(runner, session, "lro_call_id", "pendingTool", ImmutableMap.of("message", "hi")); + + // The paused leaf resumes: the answer's branch does not strand the invocation. + assertThat(simplifyEvents(resumed)).contains("leaf_agent: leaf resumed"); + } + + /** Records the user content each run is started with, so a resumed run can be checked. */ + private static final class UserContentCapturingPlugin extends BasePlugin { + private @Nullable String lastUserText; + + UserContentCapturingPlugin() { + super("user-content-capturing"); + } + + @Override + public Maybe beforeRunCallback(InvocationContext invocationContext) { + lastUserText = + invocationContext.userContent().flatMap(Content::parts).stream() + .flatMap(List::stream) + .map(part -> part.text().orElse("")) + .findFirst() + .orElse(null); + return Maybe.empty(); + } + } + + // Resuming a non-existent invocation with no new message has nothing to resume: runAsync surfaces + // IllegalArgumentException rather than running an empty model call. + @Test + public void resume_nonExistentInvocationId_throwsIllegalArgument() { + LlmAgent agent = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("x"))).name("agent").build(); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + + runner + .runAsync( + "user", + session.id(), + "does-not-exist", + /* newMessage= */ null, + RunConfig.builder().build(), + /* stateDelta= */ null) + .test() + .assertError(IllegalArgumentException.class); + } + + // With an explicit invocation id the runner skips message-based resolution entirely, so the + // orphan check has to fire on this route too. Paired with + // resume_functionResponseWithNoMatchingCall_throwsIllegalArgument, which covers the null-id + // route; a regression could plausibly skip one and not the other. + @Test + public void resume_orphanFunctionResponseWithProvidedInvocationId_throwsIllegalArgument() { + LlmAgent agent = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("x"))).name("agent").build(); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + + 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 + public void resume_withFunctionResponse_copiesBranchFromMatchingCall() { + TestBaseAgent agent = + new TestBaseAgent("test_agent", "desc", () -> Flowable.empty(), null, null, null); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + Object unusedFc = + runner + .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(); + + 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"); + } + + // A pending long-running call must stop a resumable LoopAgent after the current iteration rather + // than looping again (re-calling the model every iteration), matching Python ADK v1. + @Test + public void runAsync_loopAgentWithLongRunningSubAgent_resumable_stopsAfterFirstIteration() { + AtomicInteger calls = new AtomicInteger(); + TestLlm loopLlm = + createTestLlm( + () -> + calls.incrementAndGet() <= 5 + ? Flowable.just( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hello"))) + : Flowable.just(createTextLlmResponse("stop"))); + LlmAgent inner = + createTestAgentBuilder(loopLlm).name("inner").tools(pendingFunctionTool()).build(); + LoopAgent loop = + LoopAgent.builder() + .name("loop") + .subAgents(ImmutableList.of(inner)) + .maxIterations(3) + .build(); + Runner runner = resumableRunner(loop); + Session session = newSession(runner); + + ImmutableList unused = runTurn(runner, session, "from user"); + + // Paused after the first iteration: one model call, not maxIterations. + assertThat(loopLlm.getRequests()).hasSize(1); + } + + // In a resumable ParallelAgent, a pending long-running call pauses only its own branch (via the + // flow); other branches still complete. ParallelAgent needs no special handling, matching Python + // ADK v1 (cancelling siblings would diverge). + @Test + public void runAsync_parallelAgentWithLongRunningBranch_resumable_otherBranchCompletes() { + TestLlm longRunningLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hello")), + createTextLlmResponse("unexpected")); + LlmAgent longRunningBranch = + createTestAgentBuilder(longRunningLlm) + .name("long_running_branch") + .tools(pendingFunctionTool()) + .build(); + LlmAgent plainBranch = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("plain branch done"))) + .name("plain_branch") + .build(); + ParallelAgent parallel = + ParallelAgent.builder() + .name("parallel") + .subAgents(ImmutableList.of(longRunningBranch, plainBranch)) + .build(); + Runner runner = resumableRunner(parallel); + Session session = newSession(runner); + + ImmutableList events = runTurn(runner, session, "from user"); + + // The long-running branch paused after one model call; the other branch still completed. + assertThat(longRunningLlm.getRequests()).hasSize(1); + assertThat(simplifyEvents(events)).contains("plain_branch: plain branch done"); + } +} 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..b7adda6dd 100644 --- a/core/src/test/java/com/google/adk/runner/RunnerTest.java +++ b/core/src/test/java/com/google/adk/runner/RunnerTest.java @@ -45,8 +45,6 @@ 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.ParallelAgent; import com.google.adk.agents.RunConfig; import com.google.adk.agents.SequentialAgent; import com.google.adk.apps.App; @@ -122,14 +120,21 @@ public final class RunnerTest { @Rule public final OpenTelemetryRule openTelemetryRule = OpenTelemetryRule.create(); private final BasePlugin plugin = mockPlugin("test"); + private final Content pluginContent = createContent("from plugin"); + private final TestLlm testLlm = createTestLlm(createLlmResponse(createContent("from llm"))); + private final LlmAgent agent = createTestAgentBuilder(testLlm).build(); + private Runner runner; + private Session session; + private Tracer originalTracer; private final FailingEchoTool failingEchoTool = new FailingEchoTool(); + private final EchoTool echoTool = new EchoTool(); private final TestLlm testLlmWithFunctionCall = @@ -2353,219 +2358,6 @@ public void runAsync_withToolConfirmation_inSequentialAgentSubAgent_resumesSubAg .inOrder(); } - // OSS HITL: after an adk_request_confirmation resumes sub-agent B in a SequentialAgent(A, B, C), - // the workflow must advance to C without re-running the already completed A. - @Test - @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). - public void runAsync_withToolConfirmation_inSequentialAgent_runsLaterSubAgentsAfterResume() { - LlmAgent agentA = - createTestAgentBuilder(createTestLlm(createTextLlmResponse("agent A done"))) - .name("a_agent") - .build(); - // With resumability on, B pauses right after requesting confirmation (no extra model call), so - // a - // single follow-up response covers the resume. - TestLlm bTestLlm = - createTestLlm( - createFunctionCallLlmResponse( - "tool_call_id", "echoTool", ImmutableMap.of("message", "hello")), - createTextLlmResponse("Response after user confirmed.")); - LlmAgent agentB = - createTestAgentBuilder(bTestLlm) - .name("b_agent") - .tools(FunctionTool.create(Tools.class, "echoTool", /* requireConfirmation= */ 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 eventsBeforeConfirmation = - runner - .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) - .toList() - .blockingGet(); - - // Turn 1: A runs, B pauses for confirmation, and C must not run yet. - assertThat(simplifyEvents(eventsBeforeConfirmation)).contains("a_agent: agent A done"); - assertThat(simplifyEvents(eventsBeforeConfirmation)).doesNotContain("c_agent: agent C done"); - - FunctionCall askUserConfirmationFunctionCall = - Iterables.getOnlyElement( - eventsBeforeConfirmation.stream() - .map(Functions::getAskUserConfirmationFunctionCalls) - .filter(functionCalls -> !functionCalls.isEmpty()) - .findFirst() - .get()); - List eventsAfterConfirmation = - runner - .runAsync( - "user", - session.id(), - Content.fromParts( - Part.builder() - .functionResponse( - FunctionResponse.builder() - .id(askUserConfirmationFunctionCall.id().get()) - .name(askUserConfirmationFunctionCall.name().get()) - .response(ImmutableMap.of("confirmed", true))) - .build())) - .toList() - .blockingGet(); - - // Turn 2: B resumes and executes the tool, then C runs. A is not re-run. - assertThat(simplifyEvents(eventsAfterConfirmation)) - .containsExactly( - "b_agent: FunctionResponse(name=echoTool, response={message=hello})", - "b_agent: Response after user confirmed.", - "c_agent: agent C done") - .inOrder(); - } - - // Long-running-call HITL: a pending long-running function call (not the confirmation flow) pauses - // SequentialAgent(A, B, C) after B; on resume B continues and C runs, without re-running A. - @Test - @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). - public void runAsync_withLongRunningCall_inSequentialAgent_runsLaterSubAgentsAfterResume() { - LlmAgent agentA = - 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. - TestLlm bTestLlm = - createTestLlm( - createFunctionCallLlmResponse( - "lro_call_id", "echoTool", ImmutableMap.of("message", "hello")), - createTextLlmResponse("agent B resumed")); - LlmAgent agentB = - createTestAgentBuilder(bTestLlm) - .name("b_agent") - .tools( - FunctionTool.create( - Tools.class, - "echoTool", - /* 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("echoTool") - .response(ImmutableMap.of("message", "hello"))) - .build())) - .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(); - } - - // 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. - @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(); - 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 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"); - } - // 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. @@ -2754,109 +2546,6 @@ private static List resumeWithStatus(Runner runner, Session session, Stri .blockingGet(); } - // A pending long-running call must stop a resumable LoopAgent after the current iteration rather - // than looping again (re-calling the model every iteration), matching Python ADK v1. - @Test - @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). - public void runAsync_loopAgentWithLongRunningSubAgent_resumable_stopsAfterFirstIteration() { - AtomicInteger calls = new AtomicInteger(); - TestLlm loopLlm = - createTestLlm( - () -> - calls.incrementAndGet() <= 5 - ? Flowable.just( - createFunctionCallLlmResponse( - "lro_call_id", "echoTool", ImmutableMap.of("message", "hello"))) - : Flowable.just(createTextLlmResponse("stop"))); - LlmAgent inner = - createTestAgentBuilder(loopLlm) - .name("inner") - .tools( - FunctionTool.create( - Tools.class, - "echoTool", - /* requireConfirmation= */ false, - /* isLongRunning= */ true)) - .build(); - LoopAgent loop = - LoopAgent.builder() - .name("loop") - .subAgents(ImmutableList.of(inner)) - .maxIterations(3) - .build(); - Runner runner = - Runner.builder() - .app( - App.builder() - .name("test") - .rootAgent(loop) - .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) - .build()) - .build(); - Session session = runner.sessionService().createSession("test", "user").blockingGet(); - - List unused = - runner - .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) - .toList() - .blockingGet(); - - // Paused after the first iteration: one model call, not maxIterations. - assertThat(loopLlm.getRequests()).hasSize(1); - } - - // In a resumable ParallelAgent, a pending long-running call pauses only its own branch (via the - // flow); other branches still complete. ParallelAgent needs no special handling, matching Python - // ADK v1 (cancelling siblings would diverge). - @Test - @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). - public void runAsync_parallelAgentWithLongRunningBranch_resumable_otherBranchCompletes() { - TestLlm longRunningLlm = - createTestLlm( - createFunctionCallLlmResponse( - "lro_call_id", "echoTool", ImmutableMap.of("message", "hello")), - createTextLlmResponse("unexpected")); - LlmAgent longRunningBranch = - createTestAgentBuilder(longRunningLlm) - .name("long_running_branch") - .tools( - FunctionTool.create( - Tools.class, - "echoTool", - /* requireConfirmation= */ false, - /* isLongRunning= */ true)) - .build(); - LlmAgent plainBranch = - createTestAgentBuilder(createTestLlm(createTextLlmResponse("plain branch done"))) - .name("plain_branch") - .build(); - ParallelAgent parallel = - ParallelAgent.builder() - .name("parallel") - .subAgents(ImmutableList.of(longRunningBranch, plainBranch)) - .build(); - Runner runner = - Runner.builder() - .app( - App.builder() - .name("test") - .rootAgent(parallel) - .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 long-running branch paused after one model call; the other branch still completed. - assertThat(longRunningLlm.getRequests()).hasSize(1); - assertThat(simplifyEvents(events)).contains("plain_branch: plain branch done"); - } - // Resumability disabled (default): a SequentialAgent(A, B, C) does not pause on B's HITL call, so // all sub-agents run in the same turn — matching Python ADK v1 with resumability disabled. @Test @@ -3073,7 +2762,9 @@ public void runner_executesSaveArtifactFlow() { } private static final String BLOB_MIME_TYPE = "example/octet-stream"; + private static final String BLOB_PAYLOAD = "blob payload"; + private static final String PLACEHOLDER_FORMAT = "Uploaded file: %s. It has been saved to the artifacts"; 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..635237d74 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,64 @@ 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_agentStateNotAMap_isIgnored() { + // A session written by another ADK runtime may carry a non-map agentState. Dropping the field + // keeps the session loadable; an unchecked cast would fail the whole load. + 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", "not-a-map"); + apiEvent.put("actions", actions); + + Event event = SessionJsonConverter.fromApiEvent(apiEvent); + + assertThat(event.actions().agentState()).isEmpty(); + } + @Test public void fromApiEvent_minimalEvent_success() { Map apiEvent = new HashMap<>(); diff --git a/core/src/test/java/com/google/adk/testing/ResumabilityTestUtils.java b/core/src/test/java/com/google/adk/testing/ResumabilityTestUtils.java new file mode 100644 index 000000000..840c6ea9e --- /dev/null +++ b/core/src/test/java/com/google/adk/testing/ResumabilityTestUtils.java @@ -0,0 +1,294 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.adk.testing; + +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.simplifyResumableEvents; +import static com.google.common.collect.ImmutableList.toImmutableList; +import static com.google.common.truth.Truth.assertThat; +import static com.google.common.truth.Truth.assertWithMessage; +import static java.util.Arrays.stream; + +import com.google.adk.agents.BaseAgent; +import com.google.adk.agents.LlmAgent; +import com.google.adk.agents.RunConfig; +import com.google.adk.apps.App; +import com.google.adk.apps.ResumabilityConfig; +import com.google.adk.events.Event; +import com.google.adk.models.LlmResponse; +import com.google.adk.plugins.BasePlugin; +import com.google.adk.runner.Runner; +import com.google.adk.sessions.Session; +import com.google.adk.tools.FunctionTool; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +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 io.reactivex.rxjava3.core.Flowable; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import org.jspecify.annotations.Nullable; + +/** + * Helpers shared by the resumability tests: building a runner in each resumption mode, driving + * turns and resumes through it, and asserting on the checkpoints it emits. + * + *

Separate from {@link TestUtils} because these reach the {@link Runner} and {@link App} layer, + * which the agent-level helpers there do not. + */ +public final class ResumabilityTestUtils { + + /** Tools the resumability tests pause on. */ + public static final class Tools { + private Tools() {} + + /** Returns its argument, so a call to it completes within the turn. */ + public static ImmutableMap echoTool(String message) { + return ImmutableMap.of("message", message); + } + + /** Awaits a result supplied later, so a call to it leaves the invocation paused. */ + public static @Nullable ImmutableMap pendingTool(String message) { + return null; + } + } + + /** {@code pendingTool} wired long-running: the shape that pauses an invocation. */ + public static FunctionTool pendingFunctionTool() { + return FunctionTool.create( + Tools.class, "pendingTool", /* requireConfirmation= */ false, /* isLongRunning= */ true); + } + + /** {@code echoTool} wired long-running but returning in-turn, so it does not pause. */ + public static FunctionTool longRunningEchoFunctionTool() { + return FunctionTool.create( + Tools.class, "echoTool", /* requireConfirmation= */ false, /* isLongRunning= */ true); + } + + /** {@code echoTool} requiring confirmation: the human-in-the-loop shape. */ + public static FunctionTool confirmingEchoFunctionTool() { + return FunctionTool.create(Tools.class, "echoTool", /* requireConfirmation= */ true); + } + + /** An {@link LlmAgent} named {@code name} whose model returns {@code texts}, one per call. */ + public static LlmAgent textAgent(String name, String... texts) { + return createTestAgentBuilder( + createTestLlm( + stream(texts).map(TestUtils::createTextLlmResponse).toArray(LlmResponse[]::new))) + .name(name) + .build(); + } + + /** An {@link LlmAgent} named {@code name} carrying the long-running {@code pendingTool}. */ + public static LlmAgent pausingAgent(String name, TestLlm llm) { + return createTestAgentBuilder(llm).name(name).tools(pendingFunctionTool()).build(); + } + + /** A model script that calls {@code pendingTool} once, then says {@code then}. */ + public static TestLlm pauseThenSay(String callId, String then) { + return createTestLlm( + TestUtils.createFunctionCallLlmResponse( + callId, "pendingTool", ImmutableMap.of("message", "hello")), + createTextLlmResponse(then)); + } + + /** A {@link Runner} over {@code rootAgent} with durable resumability enabled. */ + @SuppressWarnings("deprecation") // ResumabilityConfig is @Experimental, not deprecated, here. + public static Runner resumableRunner(BaseAgent rootAgent, BasePlugin... plugins) { + return Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(rootAgent) + .plugins(ImmutableList.copyOf(plugins)) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build()) + .build(); + } + + /** A {@link Runner} over {@code rootAgent} with only the deprecated plain-text shim enabled. */ + @SuppressWarnings("deprecation") // The plain-text continuation shim is deprecated by design. + public static Runner shimRunner(BaseAgent rootAgent) { + return Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(rootAgent) + .resumabilityConfig( + ResumabilityConfig.builder().plainTextContinuationAutoResume(true).build()) + .build()) + .build(); + } + + /** A fresh session on {@code runner}'s session service. */ + public static Session newSession(Runner runner) { + return runner.sessionService().createSession("test", "user").blockingGet(); + } + + /** Runs one new plain-text turn and returns its events. */ + @CanIgnoreReturnValue + public static ImmutableList runTurn(Runner runner, Session session, String text) { + return ImmutableList.copyOf( + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText(text))) + .toList() + .blockingGet()); + } + + /** Content answering a pending function call, as a user turn carries it. */ + public static Content functionResponseContent( + String callId, String toolName, Map response) { + return Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder().id(callId).name(toolName).response(response)) + .build()); + } + + /** Resumes by answering the pending call {@code callId}; the runner resolves the invocation. */ + @CanIgnoreReturnValue + public static ImmutableList answerCall( + Runner runner, + Session session, + String callId, + String toolName, + Map response) { + return ImmutableList.copyOf( + runner + .runAsync("user", session.id(), functionResponseContent(callId, toolName, response)) + .toList() + .blockingGet()); + } + + /** The un-subscribed resume stream, for tests asserting on the error rather than the events. */ + public static Flowable resumeFlowable( + Runner runner, + Session session, + @Nullable String invocationId, + @Nullable Content newMessage, + @Nullable Map stateDelta) { + return runner.runAsync( + "user", session.id(), invocationId, newMessage, RunConfig.builder().build(), stateDelta); + } + + /** The six-arg resume overload, for tests that vary the invocation id or the state delta. */ + @CanIgnoreReturnValue + public static ImmutableList resume( + Runner runner, + Session session, + @Nullable String invocationId, + @Nullable Content newMessage, + @Nullable Map stateDelta) { + return ImmutableList.copyOf( + resumeFlowable(runner, session, invocationId, newMessage, stateDelta) + .toList() + .blockingGet()); + } + + /** Resumes {@code invocationId} with no new message. */ + @CanIgnoreReturnValue + public static ImmutableList resumeById( + Runner runner, Session session, @Nullable String invocationId) { + return resume(runner, session, invocationId, /* newMessage= */ null, /* stateDelta= */ null); + } + + /** Reloads {@code session} from the runner's session service, to assert on persisted state. */ + public static Session reloadSession(Runner runner, Session session) { + return runner + .sessionService() + .getSession("test", "user", session.id(), Optional.empty()) + .blockingGet(); + } + + /** + * A model turn issuing two long-running calls at once: the shape that pauses until both answer. + */ + public static Content twoPendingCalls(String firstId, String secondId) { + return Content.builder() + .role("model") + .parts(pendingCallPart(firstId, "a"), pendingCallPart(secondId, "b")) + .build(); + } + + private static Part pendingCallPart(String callId, String message) { + return Part.builder() + .functionCall( + FunctionCall.builder() + .id(callId) + .name("pendingTool") + .args(ImmutableMap.of("message", message)) + .build()) + .build(); + } + + /** Asserts {@code author} emitted an end-of-agent checkpoint. */ + public static void assertEndOfAgent(List events, String author) { + assertWithMessage( + "end-of-agent checkpoint for %s in %s", author, simplifyResumableEvents(events)) + .that( + events.stream() + .anyMatch(event -> Objects.equals(event.author(), author) && endsAgent(event))) + .isTrue(); + } + + /** Asserts {@code author} emitted no end-of-agent checkpoint. */ + public static void assertNoEndOfAgent(List events, String author) { + assertWithMessage( + "unexpected end-of-agent checkpoint for %s in %s", + author, simplifyResumableEvents(events)) + .that( + events.stream() + .anyMatch(event -> Objects.equals(event.author(), author) && endsAgent(event))) + .isFalse(); + } + + /** Asserts whether {@code author} emitted an agent-state checkpoint. */ + public static void assertAgentStateCheckpoint( + List events, String author, boolean expected) { + assertWithMessage( + "agent-state checkpoint for %s in %s", author, simplifyResumableEvents(events)) + .that( + events.stream() + .anyMatch( + event -> + Objects.equals(event.author(), author) + && event.actions().agentState().isPresent())) + .isEqualTo(expected); + } + + /** Asserts {@code session} holds no resumability checkpoints at all. */ + public static void assertNoCheckpoints(Session session) { + assertThat( + session.events().stream() + .filter(event -> endsAgent(event) || event.actions().agentState().isPresent()) + .collect(toImmutableList())) + .isEmpty(); + } + + private static boolean endsAgent(Event event) { + return event.actions().endOfAgent(); + } + + private ResumabilityTestUtils() {} +} 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