From 338e46ac7af75e2944deae474433a74164387a54 Mon Sep 17 00:00:00 2001 From: Kamil Tomaszek Date: Tue, 8 Sep 2026 02:22:18 -0700 Subject: [PATCH] fix(dev): let the dev UI's public address be configured Behind a gateway that strips a path prefix, the Development UI's own API calls went to the origin root, because the bundled `runtime-config.json` ships an empty `backendUrl` with no way to change it. `adk.web.backend-url` now names the address browsers reach the server on. It sets that value and prefixes the entry redirect, so both halves work without anything being forwarded by the proxy. It must be an absolute URL, because the UI reads a value without a scheme as the host of its live connection. Left unset, nothing changes. PiperOrigin-RevId: 977785275 --- .gemini/styleguide.md | 4 + AGENTS.md | 4 + CLAUDE.md | 1 + .../java/com/google/adk/agents/BaseAgent.java | 47 +- .../google/adk/agents/InvocationContext.java | 270 ++- .../java/com/google/adk/agents/LlmAgent.java | 159 +- .../java/com/google/adk/agents/LoopAgent.java | 152 +- .../com/google/adk/agents/ParallelAgent.java | 77 +- .../google/adk/agents/SequentialAgent.java | 101 +- .../adk/agents/WorkflowAgentResumption.java | 11 +- .../adk/agents/WorkflowAgentStates.java | 103 + .../main/java/com/google/adk/apps/App.java | 11 +- .../google/adk/apps/ResumabilityConfig.java | 43 +- .../adk/flows/llmflows/AgentTransfer.java | 36 +- .../adk/flows/llmflows/BaseLlmFlow.java | 159 +- .../google/adk/flows/llmflows/Functions.java | 8 +- .../google/adk/flows/llmflows/StepResume.java | 241 +++ .../java/com/google/adk/runner/Runner.java | 653 +++++- .../adk/sessions/BaseSessionService.java | 7 + .../adk/sessions/InMemorySessionService.java | 60 +- .../google/adk/sessions/SessionException.java | 3 + .../adk/sessions/SessionJsonConverter.java | 40 +- .../google/adk/tools/GoogleSearchTool.java | 25 +- .../com/google/adk/tools/UrlContextTool.java | 15 +- .../adk/agents/InvocationContextTest.java | 510 ++++- .../com/google/adk/agents/LlmAgentTest.java | 303 +++ .../com/google/adk/agents/LoopAgentTest.java | 141 ++ .../google/adk/agents/ParallelAgentTest.java | 202 ++ .../adk/agents/SequentialAgentTest.java | 166 ++ .../adk/flows/llmflows/AgentTransferTest.java | 74 + .../adk/flows/llmflows/FunctionsTest.java | 33 + .../runner/RunnerLegacyResumabilityTest.java | 565 +++-- .../adk/runner/RunnerResumabilityTest.java | 1864 +++++++++++++++++ .../sessions/InMemorySessionServiceTest.java | 87 + .../sessions/SessionJsonConverterTest.java | 101 + .../adk/testing/ResumabilityTestUtils.java | 294 +++ .../com/google/adk/testing/TestUtils.java | 40 + .../com/google/adk/tools/BaseToolTest.java | 78 +- dev/README.md | 31 + .../java/com/google/adk/web/AdkWebServer.java | 34 +- .../com/google/adk/web/config/BackendUrl.java | 124 ++ .../google/adk/web/config/DevUiAssets.java | 13 +- .../controller/RuntimeConfigController.java | 104 + .../web/AdkWebServerProxyRedirectTest.java | 12 +- .../adk/web/BackendUrlRedirectTest.java | 169 ++ .../google/adk/web/config/BackendUrlTest.java | 188 ++ .../adk/web/config/DevUiAssetsTest.java | 6 + .../RuntimeConfigControllerMergeTest.java | 123 ++ .../RuntimeConfigControllerTest.java | 121 ++ .../com/google/adk/tokt/KtRunnerToJava.kt | 36 +- .../google/adk/tokt/KtRunnerInteropTest.kt | 368 +++- 51 files changed, 7551 insertions(+), 466 deletions(-) create mode 100644 .gemini/styleguide.md create mode 100644 AGENTS.md create mode 100644 CLAUDE.md 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/runner/RunnerResumabilityTest.java create mode 100644 core/src/test/java/com/google/adk/testing/ResumabilityTestUtils.java create mode 100644 dev/src/main/java/com/google/adk/web/config/BackendUrl.java create mode 100644 dev/src/main/java/com/google/adk/web/controller/RuntimeConfigController.java create mode 100644 dev/src/test/java/com/google/adk/web/BackendUrlRedirectTest.java create mode 100644 dev/src/test/java/com/google/adk/web/config/BackendUrlTest.java create mode 100644 dev/src/test/java/com/google/adk/web/controller/RuntimeConfigControllerMergeTest.java create mode 100644 dev/src/test/java/com/google/adk/web/controller/RuntimeConfigControllerTest.java diff --git a/.gemini/styleguide.md b/.gemini/styleguide.md new file mode 100644 index 000000000..51a585547 --- /dev/null +++ b/.gemini/styleguide.md @@ -0,0 +1,4 @@ +# Gemini Code Assist Customization + +_Placeholder - customization for Gemini Code Assist in this repository is coming +soon._ diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 000000000..6d86f1b09 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,4 @@ +# AGENTS.md + +_Placeholder - guidance for AI coding agents working in this repository is coming +soon._ diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 000000000..543a9c15f --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1 @@ +See [AGENTS.md](./AGENTS.md) for project context, commands, and contribution guidelines for AI coding agents. 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 eeb04bdf1..c4bb0fbb7 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; @@ -39,6 +41,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; @@ -136,10 +139,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."); } /** @@ -468,6 +469,44 @@ 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. The + * recorded state is cleared and the agent marked finished, so a later run skips it. + * + * @param context Current invocation context. + * @return 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. The agent is left unfinished, so a later run resumes from this checkpoint. + * + * @param context Current invocation context. + * @param agentState The serialized agent state to persist. + * @return 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 384fd8372..6c874c3e1 100644 --- a/core/src/main/java/com/google/adk/agents/InvocationContext.java +++ b/core/src/main/java/com/google/adk/agents/InvocationContext.java @@ -20,9 +20,11 @@ import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.collect.ImmutableSet.toImmutableSet; +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.flows.llmflows.Functions; import com.google.adk.memory.BaseMemoryService; import com.google.adk.models.LlmCallsLimitExceededException; import com.google.adk.plugins.Plugin; @@ -31,22 +33,27 @@ 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; @@ -64,6 +71,9 @@ public class InvocationContext { private final @Nullable ResumabilityConfig resumabilityConfig; private final InvocationCostManager invocationCostManager; private final Map callbackContextData; + // Shared by reference 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; @@ -91,6 +101,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. */ @@ -304,16 +316,253 @@ public Optional contextCacheConfig() { return Optional.ofNullable(contextCacheConfig); } + /** Returns whether the current invocation is resumable. */ + @Experimental + public boolean isResumable() { + return resumabilityConfig != null && resumabilityConfig.isResumable(); + } + /** - * Returns whether the current invocation is resumable. Mirrors Python ADK v1's {@code - * InvocationContext.is_resumable}. The deprecated plain-text continuation shim selects the same - * resumption behavior, so it reports resumable too; the two are mutually exclusive. + * 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 isResumable() { - return resumabilityConfig != null - && (resumabilityConfig.isResumable() - || resumabilityConfig.isPlainTextContinuationAutoResume()); + 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) { + // LinkedHashMap, not ImmutableMap.copyOf: a deserialized agentState may carry a null value. + 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(); + // The deprecated setEndInvocation aliases endOfAgent, so only a content-less marker counts. + if (event.actions().endOfAgent() && event.content().isEmpty()) { + endOfAgents.put(author, true); + agentStates.remove(author); + } else if (agentState.isPresent()) { + setAgentState(author, agentState.get(), /* endOfAgent= */ false); + } else if (!author.equals(Role.USER) + && event.content().isPresent() + && !agentStates.containsKey(author)) { + // Content after an end-of-agent marker reopens the agent, as in Python. + 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) { + // eventsOnCurrentBranch owns branch filtering: author-asymmetric rule and cross-branch guard. + List results; + if (currentBranch) { + results = new ArrayList<>(eventsOnCurrentBranch()); + } else { + // session.events() is a synchronized list; copy it under its own monitor. + List sessionEvents = session.events(); + synchronized (sessionEvents) { + results = new ArrayList<>(sessionEvents); + } + } + if (currentInvocation) { + results.removeIf(event -> !invocationId.equals(event.invocationId())); + } + return ImmutableList.copyOf(results); + } + + /** + * Returns whether to pause the invocation right after this event: it is resumable and the event + * carries a long-running function call, including a synthetic {@code adk_request_confirmation} + * HITL request. Pausing, unlike ending, leaves the invocation resumable. + */ + @Experimental + public boolean shouldPauseInvocation(Event event) { + return isResumable() && carriesLongRunningCall(event); + } + + /** + * Returns whether the event carries a long-running function call, independent of the mode. Shared + * with the legacy flow's check so the two cannot drift apart. + */ + private static boolean carriesLongRunningCall(Event event) { + return Functions.hasPendingLongRunningCall(event); + } + + /** + * Returns whether either of the last two events on this branch pauses the invocation, the + * condition Python uses to decide whether to withhold an agent's end-of-agent checkpoint. + */ + boolean lastEventsPauseInvocation() { + ImmutableList events = events(/* currentInvocation= */ true, /* currentBranch= */ true); + return events.subList(Math.max(0, events.size() - 2), events.size()).stream() + .anyMatch(this::shouldPauseInvocation); + } + + /** + * 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 without emitting an + * event, 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); + } + + 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} 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, Set scope) { + if (events.isEmpty()) { + return false; + } + Set awaited = new HashSet<>(); + for (Event event : events) { + if (!carriesLongRunningCall(event)) { + continue; + } + if (!scope.contains(event.author())) { + continue; + } + // Only long-running calls count, as in Python: an ordinary call in the event runs inline. + Set longRunningIds = event.longRunningToolIds().orElse(ImmutableSet.of()); + for (FunctionCall call : event.functionCalls()) { + call.id().filter(longRunningIds::contains).ifPresent(awaited::add); + } + } + 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--) { + // The response event never answers itself, as in Python. + if (events.get(i).id().equals(functionResponseEvent.id())) { + continue; + } + 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 { @@ -375,6 +624,9 @@ 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, not copied: the checkpoints belong to the invocation, not a context. + this.agentStates = context.agentStates; + this.endOfAgents = context.endOfAgents; } private BaseSessionService sessionService; @@ -395,6 +647,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..684635ac9 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,135 @@ 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 sub-agent instead of re-invoking the model. + 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); + })); + } + } + // Normal path: emit an end-of-agent checkpoint so a later run can skip this agent. + Flowable events = + llmFlow.run(invocationContext).doOnNext(this::maybeSaveOutputToState); + AtomicBoolean paused = new AtomicBoolean(false); + return events + .doOnNext( + event -> { + if (invocationContext.shouldPauseInvocation(event)) { + paused.set(true); + } + }) + .concatWith( + Flowable.defer( + () -> { + // Only the checkpoint is withheld, as in Python; the flow still runs. A + // transferring agent stays open so a resume re-enters it and reaches the + // sub-agent it handed off to. + if (paused.get() || invocationContext.lastEventsPauseInvocation()) { + 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) { + // Declared targets first, as in Python, so a same-named agent on an unrelated branch is not + // resumed into; then the same tree-wide search the live transfer path uses. + for (BaseAgent target : transferTargets()) { + if (target.name().equals(targetName)) { + return target; + } + } + 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) && !lastEvent.functionResponses().isEmpty()) { + // 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(); + } + } + // A user event answering no call falls through; Python fails the resume instead. + 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 @@ -789,6 +922,30 @@ public boolean disallowTransferToPeers() { return disallowTransferToPeers; } + /** + * Returns the agents this one may transfer to: its sub-agents, plus its parent and peers unless + * disallowed. Transfer to a parent or peer needs an {@link LlmAgent} parent. Package-private + * because {@code AgentTransfer} cannot reach it and keeps its own copy of the same rule. + */ + ImmutableList transferTargets() { + List targets = new ArrayList<>(subAgents()); + BaseAgent parent = parentAgent(); + if (!(parent instanceof LlmAgent)) { + return ImmutableList.copyOf(targets); + } + if (!disallowTransferToParent()) { + targets.add(parent); + } + if (!disallowTransferToPeers()) { + for (BaseAgent peer : parent.subAgents()) { + if (!peer.name().equals(name())) { + targets.add(peer); + } + } + } + return ImmutableList.copyOf(targets); + } + public List beforeModelCallback() { return beforeModelCallback; } 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..45c7c5088 100644 --- a/core/src/main/java/com/google/adk/agents/LoopAgent.java +++ b/core/src/main/java/com/google/adk/agents/LoopAgent.java @@ -18,11 +18,14 @@ import com.google.adk.agents.ConfigAgentUtils.ConfigurationException; import com.google.adk.events.Event; +import com.google.common.collect.ImmutableMap; import com.google.errorprone.annotations.CanIgnoreReturnValue; import io.reactivex.rxjava3.core.Flowable; import java.util.List; +import java.util.Map; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; import org.jspecify.annotations.Nullable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -136,12 +139,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 +157,33 @@ 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 before each sub-agent, pause (not end) on a long-running call. + return Flowable.defer( + () -> { + Map state = invocationContext.agentStates().get(name()); + int startTimesLooped = WorkflowAgentStates.timesLooped(state); + // Unlike a sequence, a loop has no finished state to restore into, so it restarts at 0. + int resumeIndex = WorkflowAgentStates.resumeIndex(state, subAgents, logger); + int startIndex = resumeIndex == WorkflowAgentStates.NO_SUB_AGENT_NAMED ? 0 : resumeIndex; + 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 +199,125 @@ 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); + + /** Set once the loop must stop repeating, whether it closed, escalated or paused. */ + final AtomicBoolean finished = 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). + * + *

Only the previous sub-agent is checked for a silent pause, which under a {@link + * ParallelAgent} sits on a branch this agent cannot see: a call still unanswered inside the + * sub-agent about to run is what the resume is for, and must not stall it. The first sub-agent + * this run enters is re-entered unchecked, so one that pauses without emitting an event must be + * idempotent on re-entry. + */ + private Flowable runLoopIteration( + InvocationContext context, + List subAgents, + int startIndex, + LoopState state) { + AtomicInteger nextStart = new AtomicInteger(startIndex); + // Repeated rather than recursive: recursion nested one operator layer per iteration, which is + // unbounded when maxIterations is. + return Flowable.defer(() -> oneIteration(context, subAgents, nextStart.getAndSet(0), state)) + .repeatUntil(state.finished::get); + } + + /** + * Runs the sub-agents once from {@code startIndex}, marking {@code state} finished at the end. + */ + private Flowable oneIteration( + InvocationContext context, + List subAgents, + int startIndex, + LoopState state) { + return Flowable.defer( + () -> { + // Checked here so a resume already at the cap stops, not just a completed iteration. + if (maxIterations != null && state.timesLooped.get() >= maxIterations) { + state.finished.set(true); + return endOfAgentAndRecord(context); + } + AtomicReference ranLast = new AtomicReference<>(null); + Flowable iteration = + Flowable.fromIterable(subAgents.subList(startIndex, subAgents.size())) + .concatMap( + subAgent -> + Flowable.defer( + () -> { + if (state.shouldExit.get() || state.paused.get()) { + return Flowable.empty(); + } + // Catches a silent pause by the previous sub-agent. + BaseAgent previous = ranLast.get(); + if (previous != null + && context.hasUnansweredLongRunningCallIn(previous)) { + state.paused.set(true); + return Flowable.empty(); + } + ranLast.set(subAgent); + 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( + () -> { + // An escalation exit still closes the loop, as in Python. + if (state.paused.get() + || (!state.shouldExit.get() + && context.hasUnansweredLongRunningCallIn(this))) { + state.finished.set(true); + return Flowable.empty(); + } + // Anything but a pause completes the iteration, so count and clear state. + state.timesLooped.incrementAndGet(); + context.resetSubAgentStates(name()); + if (state.shouldExit.get()) { + state.finished.set(true); + return endOfAgentAndRecord(context); + } + // resuming is already false here, so the next iteration checkpoints normally. + return Flowable.empty(); + })); + }); + } + @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 d4066ed98..e78d47796 100644 --- a/core/src/main/java/com/google/adk/agents/ParallelAgent.java +++ b/core/src/main/java/com/google/adk/agents/ParallelAgent.java @@ -19,6 +19,7 @@ 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; @@ -28,6 +29,7 @@ import java.util.List; import java.util.Objects; import java.util.Set; +import java.util.concurrent.atomic.AtomicBoolean; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -184,13 +186,76 @@ protected Flowable runAsyncImpl(InvocationContext invocationContext) { .filter(Objects::nonNull) .collect(ImmutableSet.toImmutableSet()); - var updatedInvocationContext = setBranchForCurrentAgent(this, invocationContext); - List> agentFlowables = new ArrayList<>(); - for (BaseAgent subAgent : currentSubAgents) { - agentFlowables.add(subAgent.runAsync(updatedInvocationContext).subscribeOn(scheduler)); + if (!invocationContext.isResumable()) { + var updatedInvocationContext = setBranchForCurrentAgent(this, invocationContext); + List> agentFlowables = new ArrayList<>(); + for (BaseAgent subAgent : currentSubAgents) { + agentFlowables.add(subAgent.runAsync(updatedInvocationContext).subscribeOn(scheduler)); + } + return Flowable.merge(agentFlowables) + .takeUntil((Event event) -> asksThisAgentToExit(event, directSubAgentNames)); } - return Flowable.merge(agentFlowables) - .takeUntil((Event event) -> asksThisAgentToExit(event, directSubAgentNames)); + + // Resumable: skip finished branches, checkpoint the start, pause instead of ending. + 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, directSubAgentNames)) { + escalated.set(true); + } + })); + } + Flowable merged = + Flowable.merge(agentFlowables) + .takeUntil((Event event) -> asksThisAgentToExit(event, directSubAgentNames)); + + return initialCheckpoint + .concatWith(merged) + .concatWith( + Flowable.defer( + () -> { + if (paused.get()) { + return Flowable.empty(); + } + // A sub-agent that never records endOfAgent keeps this agent from ending. + boolean allEnded = + activeSubAgents.stream() + .allMatch( + a -> + invocationContext + .endOfAgents() + .getOrDefault(a.name(), false)); + if (escalated.get() || allEnded) { + return endOfAgentAndRecord(invocationContext); + } + return Flowable.empty(); + })); + }); } /** diff --git a/core/src/main/java/com/google/adk/agents/SequentialAgent.java b/core/src/main/java/com/google/adk/agents/SequentialAgent.java index 963c3d109..1ede77618 100644 --- a/core/src/main/java/com/google/adk/agents/SequentialAgent.java +++ b/core/src/main/java/com/google/adk/agents/SequentialAgent.java @@ -17,9 +17,12 @@ 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 java.util.concurrent.atomic.AtomicReference; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -90,24 +93,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 +143,76 @@ 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. + * + *

Only the previous sub-agent is checked for a silent pause: a call still unanswered inside + * the sub-agent about to run is what the resume is for, and must not stall it. The first + * sub-agent this run enters is re-entered unchecked, so one that pauses without emitting an event + * must be idempotent on re-entry. + */ + private Flowable runAsyncResumable( + InvocationContext invocationContext, List subAgents) { + // Deferred so each subscription reads the checkpoint and builds its own mutable state. + return Flowable.defer( + () -> { + Map state = invocationContext.agentStates().get(name()); + int resumeIndex = WorkflowAgentStates.resumeIndex(state, subAgents, logger); + // A checkpoint naming no sub-agent means the sequence already finished. + int startIndex = + resumeIndex == WorkflowAgentStates.NO_SUB_AGENT_NAMED + ? subAgents.size() + : resumeIndex; + AtomicBoolean paused = new AtomicBoolean(false); + AtomicBoolean resuming = new AtomicBoolean(state != null); + AtomicReference ranLast = new AtomicReference<>(null); + return Flowable.fromIterable(subAgents.subList(startIndex, subAgents.size())) + .concatMap( + subAgent -> + Flowable.defer( + () -> { + // Catches a previous sub-agent that paused without emitting an event. + BaseAgent previous = ranLast.get(); + if (paused.get() + || (previous != null + && invocationContext.hasUnansweredLongRunningCallIn( + previous))) { + paused.set(true); + return Flowable.empty(); + } + ranLast.set(subAgent); + 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..60eb936ec --- /dev/null +++ b/core/src/main/java/com/google/adk/agents/WorkflowAgentStates.java @@ -0,0 +1,103 @@ +/* + * 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 java.util.Map; +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; + } + + /** Returned by {@link #resumeIndex} when the restored checkpoint names no sub-agent. */ + static final int NO_SUB_AGENT_NAMED = -1; + + /** + * Returns the index of the sub-agent to resume from, or {@link #NO_SUB_AGENT_NAMED} when the + * checkpoint names none -- an absent or empty name, which Python treats as "already finished". + * Any other unreadable value throws rather than being read as finished, so a session written by + * another runtime cannot silently skip the whole workflow. + */ + static int resumeIndex( + @Nullable Map state, List subAgents, Logger logger) { + if (state == null) { + return 0; + } + if (!state.containsKey(CURRENT_SUB_AGENT)) { + return NO_SUB_AGENT_NAMED; + } + Object current = state.get(CURRENT_SUB_AGENT); + // Restarting the workflow would re-run side-effecting sub-agents, so fail the resume instead. + if (!(current instanceof String name)) { + throw new IllegalStateException(rejection(CURRENT_SUB_AGENT, current, "a string")); + } + return name.isEmpty() ? NO_SUB_AGENT_NAMED : findIndexForResumption(subAgents, name, logger); + } + + /** + * Returns the completed-iteration count from a Loop checkpoint, or 0 when it holds none. As with + * {@link #resumeIndex}, an unreadable value throws instead of silently restarting the count. + */ + static int timesLooped(@Nullable Map state) { + if (state == null || !state.containsKey(TIMES_LOOPED)) { + return 0; + } + Object value = state.get(TIMES_LOOPED); + if (!(value instanceof Number number)) { + throw new IllegalStateException(rejection(TIMES_LOOPED, value, "a number")); + } + return number.intValue(); + } + + /** Message for a checkpoint field that cannot be read; names the type, never the value. */ + private static String rejection(String field, @Nullable Object value, String expected) { + return String.format( + "Cannot resume: checkpoint field %s is %s, not %s.", + field, value == null ? "null" : "a " + value.getClass().getSimpleName(), expected); + } + + 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..4d2d1b3d9 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,8 @@ public Builder contextCacheConfig(ContextCacheConfig contextCacheConfig) { return this; } - /** - * Sets the app resumability config. - * - * @deprecated See {@link ResumabilityConfig}: partial feature, full resumability not yet - * available. - */ @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 590d9177a..b19e04b82 100644 --- a/core/src/main/java/com/google/adk/apps/ResumabilityConfig.java +++ b/core/src/main/java/com/google/adk/apps/ResumabilityConfig.java @@ -18,40 +18,40 @@ 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. + * Resumability config for every agent in an app: pause on a long-running call and resume from the + * last event. * - *

The two flags select the same resumption behavior and are mutually exclusive: set {@link - * #isResumable()}, or the deprecated shim, but not both. + *

The two flags select different resumption flows and are mutually exclusive: {@link + * #isResumable()} checkpoints agent state and resumes from it, {@link + * #isPlainTextContinuationAutoResume()} selects the deprecated legacy flow, and with neither set + * the app never pauses. * - * @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: resume is best-effort and at-least-once, so a resuming tool must be idempotent + * and 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. + * response -- resumes the last unfinished invocation instead of starting a new one; off by + * default. * - *

Selects the same resumption behavior as {@link #isResumable()}, with which it is mutually - * exclusive; it differs only in also resuming on a plain-text continuation. + *

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 set {@link #isResumable()} instead; + * @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 + * answer the paused call with a function response, and set {@link #isResumable()} instead; * this flag will be removed. */ @Deprecated @@ -83,16 +83,15 @@ public abstract static class Builder { /** * Builds the config, rejecting a combination of flags that has no defined behavior. * - * @throws IllegalArgumentException if both resumability and the deprecated shim are set; they - * select the same behavior, so exactly one may be enabled. + * @throws IllegalArgumentException if both resumability and the legacy shim are enabled; they + * select different resumption flows, so exactly one may be set. */ - @SuppressWarnings("deprecation") // Validating the deprecated shim against the supported flag. public ResumabilityConfig build() { ResumabilityConfig config = autoBuild(); checkArgument( !(config.isResumable() && config.isPlainTextContinuationAutoResume()), "resumable and plainTextContinuationAutoResume are mutually exclusive: set resumable for" - + " the supported flag, or the deprecated shim, but not both."); + + " 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/AgentTransfer.java b/core/src/main/java/com/google/adk/flows/llmflows/AgentTransfer.java index 0a0da8761..62ff664b5 100644 --- a/core/src/main/java/com/google/adk/flows/llmflows/AgentTransfer.java +++ b/core/src/main/java/com/google/adk/flows/llmflows/AgentTransfer.java @@ -44,16 +44,15 @@ public Single processRequest( "Base agent in InvocationContext is not an instance of Agent."); } - List transferTargets = getTransferTargets(agent); - if (transferTargets.isEmpty()) { + List targets = transferTargets(agent); + if (targets.isEmpty()) { return Single.just( RequestProcessor.RequestProcessingResult.create(request, ImmutableList.of())); } LlmRequest.Builder builder = request.toBuilder() - .appendInstructions( - ImmutableList.of(buildTargetAgentsInstructions(agent, transferTargets))); + .appendInstructions(ImmutableList.of(buildTargetAgentsInstructions(agent, targets))); FunctionTool agentTransferTool = createTransferToAgentTool(); agentTransferTool.processLlmRequest(builder, ToolContext.builder(context).build()); @@ -119,30 +118,29 @@ private String buildTargetAgentsInstructions(LlmAgent agent, List tra return sb.toString(); } - /** Returns valid transfer targets: sub-agents, parent, and peers (if allowed). */ - private List getTransferTargets(LlmAgent agent) { - List transferTargets = new ArrayList<>(); - transferTargets.addAll(agent.subAgents()); // Add all sub-agents - + /** + * Returns the agents {@code agent} may transfer to: its sub-agents, plus its parent and peers + * unless disallowed. Transfer to a parent or peer needs an {@link LlmAgent} parent. Order is + * load-bearing: it is the order the target list is rendered to the model in. {@code LlmAgent} + * keeps the same rule for its own use, since it cannot reach this one. + */ + static ImmutableList transferTargets(LlmAgent agent) { + List targets = new ArrayList<>(agent.subAgents()); BaseAgent parent = agent.parentAgent(); - // Agents eligible to transfer must have an LLM-based agent parent. if (!(parent instanceof LlmAgent)) { - return transferTargets; + return ImmutableList.copyOf(targets); } - if (!agent.disallowTransferToParent()) { - transferTargets.add(parent); + targets.add(parent); } - if (!agent.disallowTransferToPeers()) { - for (BaseAgent peerAgent : parent.subAgents()) { - if (!peerAgent.name().equals(agent.name())) { - transferTargets.add(peerAgent); + for (BaseAgent peer : parent.subAgents()) { + if (!peer.name().equals(agent.name())) { + targets.add(peer); } } } - - return transferTargets; + return ImmutableList.copyOf(targets); } @Schema( 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..f255bbe03 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,30 @@ private Flowable runOneStep(Context spanContext, InvocationContext contex logger.debug("End invocation requested during preprocessing."); return Flowable.empty(); } + // A tool answered during preprocessing would be replayed below. + if (context.isResumable() && preprocessAnswered.get()) { + logger.debug("Preprocessing produced a final response; ending the step."); + return Flowable.empty(); + } + + // As in Python, decide before the model call: the branch may owe work. + 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."); + // Re-id as Python's replay does; PersistBarrier skips an id-less event. + return runFunctionCalls( + context, + resume.replayEvent(), + llmRequestAfterPreprocess, + spanContext) + .map(event -> event.toBuilder().id(Event.generateEventId()).build()) + .concatMap(event -> followTransfer(event, context, spanContext)); + } try { context.incrementLlmCallsCount(); @@ -477,27 +509,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 +533,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 +548,8 @@ 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)) { + // 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 +769,86 @@ 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))); + } + BaseAgent target = nextAgent.get(); + // Resumable only: restricting the default flow too would move behavior for existing callers. + if (context.isResumable() && context.agent() instanceof LlmAgent llmAgent) { + Optional declared = declaredTransferTarget(llmAgent, agentToTransfer); + if (declared.isEmpty()) { + return self.concatWith( + Flowable.error( + new IllegalStateException( + "Agent " + + llmAgent.name() + + " is not allowed to transfer to the requested agent."))); + } + target = declared.get(); + } + return self.concatWith(target.runAsync(context).compose(Tracing.withContext(spanContext))); + } + + /** + * Returns the agent {@code from} may transfer to under {@code name}, or empty when the transfer + * is not allowed. Resolves through the declared targets rather than the tree, as Python does, so + * a same-named agent on an unrelated branch cannot stand in for the declared one. + */ + private static Optional declaredTransferTarget(LlmAgent from, String name) { + if (name.equals(from.name())) { + return Optional.of(from); + } + return AgentTransfer.transferTargets(from).stream() + .filter(candidate -> candidate.name().equals(name)) + .findFirst(); + } + + /** + * 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 3a1d45c2f..4a04cb94a 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 @@ -457,9 +457,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..35acd4ee8 --- /dev/null +++ b/core/src/main/java/com/google/adk/flows/llmflows/StepResume.java @@ -0,0 +1,241 @@ +/* + * 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; + } + } + // Calls on the last event are unanswered by construction -- nothing follows them to answer. + Event last = events.get(events.size() - 1); + if (!last.partial().orElse(false) + && !last.functionCalls().isEmpty() + // A caller-supplied event would otherwise run a tool with no model turn behind it. + && 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<>(); + // An id-less call can never be answered, so it pauses; Python keeps a None in its id set. + boolean hasIdlessCall = false; + for (FunctionCall call : callEvent.functionCalls()) { + call.name().ifPresent(callNames::add); + call.id().ifPresent(callIds::add); + hasIdlessCall |= call.id().isEmpty(); + } + 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() || hasIdlessCall) + && 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; + } + // Union, as in Python: every call on a pausing event plus its long-running ids. A different + // rule from shouldPauseInvocation's, which intersects the two. + 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); + } + } + // 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 2642033a3..7ec2e9265 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,11 @@ import com.google.adk.agents.InvocationContext; import com.google.adk.agents.LiveRequestQueue; import com.google.adk.agents.LlmAgent; +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 +61,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 +76,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 +369,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 +418,7 @@ private Single createUserMessageEvent( .id(Event.generateEventId()) .invocationId(invocationContext.invocationId()) .author(Role.USER) + .branch(branch) .content(messageToAppend); // Add state delta if provided @@ -450,6 +476,104 @@ public Flowable runAsync( this.runAsyncImpl(session, newMessage, runConfig, stateDelta))) .compose(Tracing.trace("invocation")); + return serializePerSession(sessionId, result); + } + + /** See {@link #runAsync(String, String, Content, RunConfig, Map)}. */ + public Flowable runAsync( + SessionKey sessionKey, + Content newMessage, + RunConfig runConfig, + @Nullable Map stateDelta) { + return runAsync(sessionKey.userId(), sessionKey.id(), newMessage, runConfig, stateDelta); + } + + /** See {@link #runAsync(String, String, Content, RunConfig, Map)}. */ + public Flowable runAsync(SessionKey sessionKey, Content newMessage, RunConfig runConfig) { + return runAsync(sessionKey, newMessage, runConfig, /* stateDelta= */ null); + } + + /** See {@link #runAsync(String, String, Content, RunConfig, Map)}. */ + public Flowable runAsync(SessionKey sessionKey, Content newMessage) { + return runAsync(sessionKey, newMessage, RunConfig.builder().build()); + } + + /** See {@link #runAsync(String, String, Content, RunConfig, Map)}. */ + public Flowable runAsync(String userId, String sessionId, Content newMessage) { + 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."); + Flowable result = + Flowable.defer( + () -> + this.sessionService + .getSession(appName, userId, sessionId, Optional.empty()) + .switchIfEmpty( + Single.defer( + () -> { + // Same get-or-create as a new run, as in Python. + if (runConfig.autoCreateSession()) { + return this.sessionService.createSession( + appName, userId, (Map) null, sessionId); + } + return 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")); + + // Resuming reads session events around an append, so it shares the new run's per-session turn. + return serializePerSession(sessionId, result); + } + + /** Returns {@code content} with the user role applied when it carries none, as Python does. */ + private static @Nullable Content withUserRole(@Nullable Content content) { + if (content == null || content.role().filter(role -> !role.isEmpty()).isPresent()) { + return content; + } + return content.toBuilder().role(Role.USER).build(); + } + + /** + * Chains {@code result} behind any run already in flight for {@code sessionId}, so two runs never + * read and append the session's events concurrently. A null session id runs unsequenced. + */ + private Flowable serializePerSession(@Nullable String sessionId, Flowable result) { return Flowable.defer( () -> { if (sessionId == null) { @@ -480,30 +604,6 @@ public Flowable runAsync( }); } - /** See {@link #runAsync(String, String, Content, RunConfig, Map)}. */ - public Flowable runAsync( - SessionKey sessionKey, - Content newMessage, - RunConfig runConfig, - @Nullable Map stateDelta) { - return runAsync(sessionKey.userId(), sessionKey.id(), newMessage, runConfig, stateDelta); - } - - /** See {@link #runAsync(String, String, Content, RunConfig, Map)}. */ - public Flowable runAsync(SessionKey sessionKey, Content newMessage, RunConfig runConfig) { - return runAsync(sessionKey, newMessage, runConfig, /* stateDelta= */ null); - } - - /** See {@link #runAsync(String, String, Content, RunConfig, Map)}. */ - public Flowable runAsync(SessionKey sessionKey, Content newMessage) { - return runAsync(sessionKey, newMessage, RunConfig.builder().build()); - } - - /** See {@link #runAsync(String, String, Content, RunConfig, Map)}. */ - public Flowable runAsync(String userId, String sessionId, Content newMessage) { - return runAsync(userId, sessionId, newMessage, RunConfig.builder().build()); - } - /** * Runs the agent asynchronously using a provided Session object. * @@ -521,6 +621,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 +671,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 +729,35 @@ 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 persist / plugin-callback pipeline shared by both + * invocation paths, through the after-run and compaction brackets. A {@code beforeRunEvent} + * short-circuits the agent run; {@code onEachPersisted} runs after each persisted event. + */ + 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 +767,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 +803,367 @@ private Completable compactEvents(Session session) { .orElseGet(Completable::complete); } + /** Wraps before-run callback content as a model event that short-circuits the agent run. */ + 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 rawMessage}; otherwise starts a new invocation. Requires + * resumability. + */ + private Flowable runResumableFromSession( + Session session, + @Nullable String providedInvocationId, + @Nullable Content rawMessage, + RunConfig runConfig, + @Nullable Map stateDelta) { + // As in Python, a role-less message is a user turn; otherwise the request drops it. + Content newMessage = withUserRole(rawMessage); + return Flowable.defer( + () -> { + // Rejected before the invocation-id fallback; defer turns the throw 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 (session.events().isEmpty()) { + // Only a wholly empty session is rejected; Python does not reject an unmatched id. + return Flowable.error( + new IllegalArgumentException("Session has no events to resume.")); + } + 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(); + + // Not new input: Python drops it and its state delta rather than adding a second prompt. + boolean supersededByOriginalMessage = + newMessage != null + && !carriesFunctionResponse(newMessage) + && findOriginalUserMessage(session, resolvedInvocationId).isPresent(); + + // Merged eagerly so the on-user-message callback sees it. + if (!supersededByOriginalMessage && stateDelta != null && !stateDelta.isEmpty()) { + stateDelta.forEach((key, value) -> session.state().put(key, value)); + } + + // For the pre-run plugin callbacks; the agent is resolved later, in runResumedAgent. + InvocationContext initialContext = + newInvocationContextBuilder(session) + .invocationId(resolvedInvocationId) + .runConfig(runConfig) + .userContent( + newMessage != null + ? newMessage + : originalUserMessage(session, resolvedInvocationId)) + .build(); + + Flowable events; + if (supersededByOriginalMessage) { + events = + runResumedAgent(session, resolvedInvocationId, /* userContent= */ null, runConfig); + } else if (newMessage != null) { + // Inherit the branch of the call it answers, so routing and rehydration see it. + String branch = + matchingFunctionCallEvent(session, newMessage).flatMap(Event::branch).orElse(null); + events = + this.pluginManager + .onUserMessageCallback(initialContext, newMessage) + .compose(Tracing.withContext(capturedContext)) + .defaultIfEmpty(newMessage) + .flatMap( + content -> + appendNewMessageToSession( + session, + content, + initialContext, + runConfig.saveInputBlobsAsArtifacts(), + stateDelta, + branch)) + // Persist first: 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 a content-less event, as Python does. + 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( + () -> { + BaseAgent rootAgent = this.agent; + InvocationContext rootContext = + newInvocationContextBuilder(session, rootAgent) + .invocationId(resolvedInvocationId) + .runConfig(runConfig) + .userContent( + userContent != null + ? userContent + : originalUserMessage(session, resolvedInvocationId)) + .build(); + // Rehydrate before routing, as Python does, so a root that checkpointed re-enters itself. + rootContext.populateInvocationAgentStates(); + BaseAgent resumeAgent = + rootContext.endOfAgents().containsKey(rootAgent.name()) + ? rootAgent + : findAgentToRun(session, rootAgent); + InvocationContext context = + resumeAgent.equals(rootAgent) + ? rootContext + : rootContext.toBuilder() + .agent(resumeAgent) + .branch(resumeParentBranch(session, resolvedInvocationId, resumeAgent)) + .build(); + + // 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: one whose first 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. Text alongside a function response is + * accepted, as on Python's agent path. + * + * @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; + } + // Python checks the first response only, and raises rather than opening a new invocation that + // would carry an orphan response. + checkArgument( + responses.get(0).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()); + } + + 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 findOriginalUserMessage(session, invocationId) + .orElseThrow( + () -> + new IllegalArgumentException( + "No user message to resume the requested invocation from.")); + } + + private static boolean carriesFunctionResponse(Content message) { + return message.parts().stream() + .flatMap(List::stream) + .anyMatch(part -> part.functionResponse().isPresent()); + } + + /** The user message that started {@code invocationId}, or empty when it has none. */ + private static Optional findOriginalUserMessage(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(); + } + private void copySessionStates(Session source, Session target) { // TODO: remove this hack when deprecating all runAsync with Session. target.state().putAll(source.state()); @@ -725,7 +1195,7 @@ private InvocationContext newInvocationContextForLive( } } InvocationContext.Builder builder = - newInvocationContextBuilder(session) + newInvocationContextBuilder(session, findAgentToRun(session, this.agent)) .runConfig(runConfigBuilder.build()) .userContent(Content.fromParts()) .liveRequestQueue(liveRequestQueue); @@ -734,18 +1204,21 @@ private InvocationContext newInvocationContextForLive( } private InvocationContext.Builder newInvocationContextBuilder(Session session) { - BaseAgent rootAgent = this.agent; + return newInvocationContextBuilder(session, this.findAgentToRun(session, this.agent)); + } + + /** Overload for callers that already resolved the agent, so the walk is not repeated. */ + private InvocationContext.Builder newInvocationContextBuilder(Session session, BaseAgent agent) { return InvocationContext.builder() .sessionService(this.sessionService) .artifactService(this.artifactService) .memoryService(this.memoryService) .pluginManager(this.pluginManager) - .agent(rootAgent) .session(session) .eventsCompactionConfig(this.eventsCompactionConfig) .contextCacheConfig(this.contextCacheConfig) .resumabilityConfig(this.resumabilityConfig) - .agent(this.findAgentToRun(session, rootAgent)); + .agent(agent); } public Flowable runLive( @@ -832,10 +1305,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)); }); @@ -863,30 +1333,34 @@ private boolean isTransferableAcrossAgentTree(BaseAgent agentToRun) { return true; } + /** Returns whether resumability is enabled for this runner's app. */ + private boolean isResumable() { + return resumabilityConfig != null && resumabilityConfig.isResumable(); + } + /** - * Returns whether resumability is enabled for this runner's app, by either the supported flag or - * the deprecated plain-text continuation shim, which selects the same behavior. + * 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 isResumable() { - return resumabilityConfig != null - && (resumabilityConfig.isResumable() - || resumabilityConfig.isPlainTextContinuationAutoResume()); + 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. + // Route a function response to its call's author, as Python's _agent_router does. An enclosing + // workflow advances by re-entering the root, which holds its own checkpoint, not from here. 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 (isLegacyResumability()) { + return topmostSequentialAncestor(author); + } + return author; } List events = new ArrayList<>(session.events()); @@ -901,6 +1375,12 @@ private BaseAgent findAgentToRun(Session session, BaseAgent rootAgent) { continue; } + // Markers carry no model turn; resumable-only, since endOfAgent shares endInvocation's field. + if (isResumable() + && (event.actions().endOfAgent() || event.actions().agentState().isPresent())) { + continue; + } + if (author.equals(rootAgent.name())) { return rootAgent; } @@ -921,9 +1401,8 @@ private BaseAgent findAgentToRun(Session session, BaseAgent rootAgent) { /** * Returns the top-most ancestor reachable from {@code agent} through {@link SequentialAgent} - * parents, or {@code agent} itself otherwise. Only SequentialAgent is resume-aware; other - * workflow agents are left to resume their paused sub-agent directly (via the function-call - * author). + * parents, or {@code agent} itself otherwise. Used by the deprecated legacy flow, whose only + * resume-aware workflow agent is SequentialAgent. */ private static BaseAgent topmostSequentialAncestor(BaseAgent agent) { BaseAgent result = agent; diff --git a/core/src/main/java/com/google/adk/sessions/BaseSessionService.java b/core/src/main/java/com/google/adk/sessions/BaseSessionService.java index 8596f8eb6..655289014 100644 --- a/core/src/main/java/com/google/adk/sessions/BaseSessionService.java +++ b/core/src/main/java/com/google/adk/sessions/BaseSessionService.java @@ -35,6 +35,10 @@ * methods for creating, retrieving, listing, and deleting sessions, as well as listing and * appending events to a session. Implementations of this interface handle the underlying storage * and retrieval logic. + * + *

Implementations differ on whether they reject a duplicate session id; those that do signal it + * with a {@link SessionException} whose message is {@link SessionException#SESSION_ALREADY_EXISTS}. + * This applies to every {@code createSession} overload that accepts a session id. */ public interface BaseSessionService { @@ -68,6 +72,9 @@ Single createSession( * @param sessionId An optional client-provided identifier for the session. If empty or null, the * service should generate a unique ID. * @return The newly created {@link Session} instance. + * @throws SessionException with {@link SessionException#SESSION_ALREADY_EXISTS} if {@code + * sessionId} is already in use for this app and user; only implementations that reject + * duplicates throw it. * @throws SessionException if creation fails. */ default Single createSession( diff --git a/core/src/main/java/com/google/adk/sessions/InMemorySessionService.java b/core/src/main/java/com/google/adk/sessions/InMemorySessionService.java index e54289b76..e9c925cc1 100644 --- a/core/src/main/java/com/google/adk/sessions/InMemorySessionService.java +++ b/core/src/main/java/com/google/adk/sessions/InMemorySessionService.java @@ -36,6 +36,8 @@ import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; import org.jspecify.annotations.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** * An in-memory implementation of {@link BaseSessionService} assuming {@link Session} objects are @@ -49,6 +51,24 @@ * during retrieval operations ({@code getSession}, {@code createSession}). */ public final class InMemorySessionService implements BaseSessionService { + + private static final Logger logger = LoggerFactory.getLogger(InMemorySessionService.class); + + /** What {@code createSession} does when {@code sessionId} is already in use. */ + public enum DuplicateSessionIdBehavior { + /** + * Replaces the stored session, discarding its events and state. + * + * @deprecated Construct the service with {@link #REJECT} instead, which becomes the default in + * a future release. + */ + @Deprecated + OVERWRITE, + + /** Rejects the call with {@link SessionException}. Becomes the default in a future release. */ + REJECT; + } + // Structure: appName -> userId -> sessionId -> Session private final ConcurrentMap>> sessions; @@ -58,8 +78,26 @@ public final class InMemorySessionService implements BaseSessionService { // Structure: appName -> stateKey -> stateValue private final ConcurrentMap> appState; - /** Creates a new instance of the in-memory session service with empty storage. */ + private final DuplicateSessionIdBehavior duplicateSessionIdBehavior; + + /** + * Creates a new instance of the in-memory session service with empty storage. A duplicate session + * ID overwrites the stored session, which stays the default until callers have migrated. + */ public InMemorySessionService() { + this(DuplicateSessionIdBehavior.OVERWRITE); + } + + /** + * Creates a new instance of the in-memory session service with empty storage. + * + * @param duplicateSessionIdBehavior What {@code createSession} does when the session ID is + * already in use. + */ + public InMemorySessionService(DuplicateSessionIdBehavior duplicateSessionIdBehavior) { + this.duplicateSessionIdBehavior = + Objects.requireNonNull( + duplicateSessionIdBehavior, "duplicateSessionIdBehavior cannot be null"); this.sessions = new ConcurrentHashMap<>(); this.userState = new ConcurrentHashMap<>(); this.appState = new ConcurrentHashMap<>(); @@ -102,10 +140,22 @@ public Single createSession( .lastUpdateTime(Instant.now()) .build(); - sessions - .computeIfAbsent(appName, unused -> new ConcurrentHashMap<>()) - .computeIfAbsent(userId, unused -> new ConcurrentHashMap<>()) - .put(resolvedSessionId, newSession); + ConcurrentMap userSessions = + sessions + .computeIfAbsent(appName, unused -> new ConcurrentHashMap<>()) + .computeIfAbsent(userId, unused -> new ConcurrentHashMap<>()); + if (duplicateSessionIdBehavior == DuplicateSessionIdBehavior.REJECT) { + if (userSessions.putIfAbsent(resolvedSessionId, newSession) != null) { + return Single.error(new SessionException(SessionException.SESSION_ALREADY_EXISTS)); + } + } else { + Session previous = userSessions.put(resolvedSessionId, newSession); + if (previous != null) { + logger.info( + "Replaced an existing session because its ID was reused; its events and state are" + + " discarded."); + } + } // Create a mutable copy for the return value Session returnCopy = copySession(newSession); diff --git a/core/src/main/java/com/google/adk/sessions/SessionException.java b/core/src/main/java/com/google/adk/sessions/SessionException.java index 514811187..bf0b2bda5 100644 --- a/core/src/main/java/com/google/adk/sessions/SessionException.java +++ b/core/src/main/java/com/google/adk/sessions/SessionException.java @@ -19,6 +19,9 @@ /** Represents a general error that occurred during session management operations. */ public class SessionException extends RuntimeException { + /** Message used when {@code createSession} is given a session id that is already taken. */ + public static final String SESSION_ALREADY_EXISTS = "Session already exists"; + public SessionException(String message) { super(message); } 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..1e6c5327a 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))); @@ -151,6 +152,15 @@ private static Content convertMapToContent(Object rawContentValue) { } } + /** Returns {@code key} from {@code actions}, falling back to the copy under {@code rawEvent}. */ + private static @Nullable Object checkpointValue( + @Nullable Map actions, @Nullable Map rawActions, String key) { + if (actions != null && actions.get(key) != null) { + return actions.get(key); + } + return rawActions == null ? null : rawActions.get(key); + } + /** * Converts raw API event data into an {@link Event} object. * @@ -180,10 +190,6 @@ static Event fromApiEvent(Map apiEvent) { if (escalate != null) { eventActionsBuilder.escalate(escalate); } - Boolean endOfAgent = (Boolean) actionsMap.get("endOfAgent"); - if (endOfAgent != null) { - eventActionsBuilder.endOfAgent(endOfAgent); - } eventActionsBuilder.requestedAuthConfigs( Optional.ofNullable(actionsMap.get("requestedAuthConfigs")) .map(SessionJsonConverter::asConcurrentMapOfConcurrentMaps) @@ -194,6 +200,32 @@ static Event fromApiEvent(Map apiEvent) { .orElse(new ConcurrentHashMap<>())); } + // Python's Vertex session service leaves agentState and endOfAgent out of "actions" and keeps + // them only under "rawEvent", so a session it wrote would otherwise load with no checkpoints. + Map rawActions = null; + Object rawEvent = apiEvent.get("rawEvent"); + if (rawEvent instanceof Map) { + Object actions = ((Map) rawEvent).get("actions"); + if (actions instanceof Map) { + rawActions = (Map) actions; + } + } + // Resolved per key, not as a whole map: a rawEvent carrying only one of the two must not + // shadow the other where "actions" does carry it. + Object endOfAgent = checkpointValue(actionsMap, rawActions, "endOfAgent"); + if (endOfAgent instanceof Boolean value) { + eventActionsBuilder.endOfAgent(value); + } else if (endOfAgent != null) { + logger.warn("Ignoring 'endOfAgent' of unexpected type {}", endOfAgent.getClass().getName()); + } + Object agentState = checkpointValue(actionsMap, rawActions, "agentState"); + if (agentState instanceof Map) { + eventActionsBuilder.agentState((Map) agentState); + } else if (agentState != null) { + // Drop a foreign agentState rather than fail the whole session load on a bad cast. + logger.warn("Ignoring 'agentState' of unexpected type {}", agentState.getClass().getName()); + } + Event event = Event.builder() .id((String) Iterables.getLast(Splitter.on('/').split(apiEvent.get("name").toString()))) diff --git a/core/src/main/java/com/google/adk/tools/GoogleSearchTool.java b/core/src/main/java/com/google/adk/tools/GoogleSearchTool.java index d5ff9492c..3ce5803b4 100644 --- a/core/src/main/java/com/google/adk/tools/GoogleSearchTool.java +++ b/core/src/main/java/com/google/adk/tools/GoogleSearchTool.java @@ -25,8 +25,8 @@ import java.util.List; /** - * A built-in tool that is automatically invoked by Gemini 2 and 3 models to retrieve search results - * from Google Search. + * A built-in tool that is automatically invoked by Gemini models to retrieve search results from + * Google Search. * *

This tool operates internally within the model and does not require or perform local code * execution. @@ -60,27 +60,10 @@ public Completable processLlmRequest( List existingTools = configBuilder.build().tools().orElse(ImmutableList.of()); ImmutableList.Builder updatedToolsBuilder = ImmutableList.builder(); updatedToolsBuilder.addAll(existingTools); - - String model = llmRequestBuilder.build().model().orElse(null); - if (isSupportedModel(model)) { - - updatedToolsBuilder.add(Tool.builder().googleSearch(GoogleSearch.builder().build()).build()); - configBuilder.tools(updatedToolsBuilder.build()); - } else { - return Completable.error( - new IllegalArgumentException("Google search tool is not supported for model " + model)); - } + updatedToolsBuilder.add(Tool.builder().googleSearch(GoogleSearch.builder().build()).build()); + configBuilder.tools(updatedToolsBuilder.build()); llmRequestBuilder.config(configBuilder.build()); return Completable.complete(); } - - private boolean isSupportedModel(String model) { - if (model == null || !model.startsWith("gemini-")) { - return false; - } - return model.startsWith("gemini-2") - || model.startsWith("gemini-3") - || model.endsWith("-latest"); - } } diff --git a/core/src/main/java/com/google/adk/tools/UrlContextTool.java b/core/src/main/java/com/google/adk/tools/UrlContextTool.java index fe7f9c77e..acf66dd7d 100644 --- a/core/src/main/java/com/google/adk/tools/UrlContextTool.java +++ b/core/src/main/java/com/google/adk/tools/UrlContextTool.java @@ -25,8 +25,8 @@ import java.util.List; /** - * A built-in tool that is automatically invoked by Gemini 2 and 3 models to retrieve information - * from the given URLs. + * A built-in tool that is automatically invoked by Gemini models to retrieve information from the + * given URLs. * *

This tool operates internally within the model and does not require or perform local code * execution. @@ -60,15 +60,8 @@ public Completable processLlmRequest( List existingTools = configBuilder.build().tools().orElse(ImmutableList.of()); ImmutableList.Builder updatedToolsBuilder = ImmutableList.builder(); updatedToolsBuilder.addAll(existingTools); - - String model = llmRequestBuilder.build().model().get(); - if (model != null && (model.startsWith("gemini-2") || model.startsWith("gemini-3"))) { - updatedToolsBuilder.add(Tool.builder().urlContext(UrlContext.builder().build()).build()); - configBuilder.tools(updatedToolsBuilder.build()); - } else { - return Completable.error( - new IllegalArgumentException("Url context tool is not supported for model " + model)); - } + updatedToolsBuilder.add(Tool.builder().urlContext(UrlContext.builder().build()).build()); + configBuilder.tools(updatedToolsBuilder.build()); llmRequestBuilder.config(configBuilder.build()); return Completable.complete(); 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 995387157..0bf6593cc 100644 --- a/core/src/test/java/com/google/adk/agents/InvocationContextTest.java +++ b/core/src/test/java/com/google/adk/agents/InvocationContextTest.java @@ -23,6 +23,7 @@ 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; @@ -31,6 +32,7 @@ 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; @@ -74,30 +76,30 @@ public void setUp() { activeStreamingTools.put("test-tool", new ActiveStreamingTool(new LiveRequestQueue())); } - // The deprecated shim selects the same resumption behavior as resumable(true), so every - // resumability branch keyed on isResumable() must treat it identically. + // The two flags select different flows: the shim runs the legacy one and is not resumable. @Test - public void isResumable_shimOnly_reportsResumable() { - InvocationContext shimContext = contextWith(resumabilityConfigWithShim()); - InvocationContext resumableContext = contextWith(resumabilityConfigResumable()); + @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()).isTrue(); - assertThat(resumableContext.isResumable()).isTrue(); - assertThat(neitherContext.isResumable()).isFalse(); - } + assertThat(shimContext.isResumable()).isFalse(); + assertThat(shimContext.isLegacyResumability()).isTrue(); - @SuppressWarnings("deprecation") // Exercises the deprecated shim. - private static ResumabilityConfig resumabilityConfigWithShim() { - return ResumabilityConfig.builder().plainTextContinuationAutoResume(true).build(); - } + assertThat(resumableContext.isResumable()).isTrue(); + assertThat(resumableContext.isLegacyResumability()).isFalse(); - @SuppressWarnings("deprecation") // ResumabilityConfig is deprecated until durable resumability. - private static ResumabilityConfig resumabilityConfigResumable() { - return ResumabilityConfig.builder().resumable(true).build(); + // 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(); } - @SuppressWarnings("deprecation") // ResumabilityConfig is deprecated until durable resumability. private InvocationContext contextWith(ResumabilityConfig resumabilityConfig) { return InvocationContext.builder() .sessionService(mockSessionService) @@ -921,6 +923,19 @@ private static Event agentEvent(@Nullable String branch) { return Event.builder().author("some_agent").branch(branch).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(); + } + private static Event callEvent(@Nullable String branch, String callId) { return Event.builder() .author("some_agent") @@ -949,4 +964,465 @@ private static Event userResponseEvent(@Nullable String branch, String callId) { .build())) .build(); } + + // ---- 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(); + } + + @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 longRunningCallEvent(String callId, String name) { + return Event.builder() + .id("m-" + callId) + .invocationId("inv") + .author("root") + .content( + Content.fromParts( + Part.builder() + .functionCall(FunctionCall.builder().id(callId).name(name).build()) + .build())) + .longRunningToolIds(ImmutableSet.of(callId)) + .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 lastEventsPauseInvocation_callInsideWindow_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.lastEventsPauseInvocation()).isTrue(); + } + + @Test + public void lastEventsPauseInvocation_callPushedOutOfWindow_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.lastEventsPauseInvocation()).isFalse(); + } + + // Python looks only at whether the event carries a long-running call, never at whether it was + // answered, so a fully answered call still inside the window pauses. + @Test + public void lastEventsPauseInvocation_answeredCallInsideWindow_returnsTrue() { + Session eventSession = Session.builder("s").build(); + eventSession.events().add(longRunningCallEvent("c1", "approve_c1")); + eventSession.events().add(functionResponseEvent("c1", "approve_c1")); + InvocationContext context = resumableContext(eventSession, "inv"); + + assertThat(context.lastEventsPauseInvocation()).isTrue(); + } + + // The converse: an unanswered call older than the two-event window does not pause, as in Python. + @Test + public void lastEventsPauseInvocation_unansweredCallOlderThanWindow_returnsFalse() { + Session eventSession = Session.builder("s").build(); + eventSession.events().add(twoLongRunningCallsEvent()); + eventSession + .events() + .add(agentEvent("inv", "root", null, Content.fromParts(Part.fromText("thinking")))); + eventSession + .events() + .add(agentEvent("inv", "root", null, Content.fromParts(Part.fromText("still here")))); + InvocationContext context = resumableContext(eventSession, "inv"); + + assertThat(context.lastEventsPauseInvocation()).isFalse(); + } + + @Test + public void lastEventsPauseInvocation_noCall_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.lastEventsPauseInvocation()).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= */ null)); + InvocationContext context = resumableContext(eventSession, "inv"); + + context.populateInvocationAgentStates(); + + assertThat(context.endOfAgents()).containsEntry("a", true); + assertThat(context.agentStates()).doesNotContainKey("a"); + } + + // EventActions.setEndInvocation aliases endOfAgent, so only a content-less event counts as done. + @Test + public void populateInvocationAgentStates_endOfAgentOnContentEvent_notTreatedAsCheckpoint() { + Session eventSession = Session.builder("s").build(); + eventSession + .events() + .add( + agentEvent( + "inv", + "a", + EventActions.builder().endOfAgent(true).build(), + Content.fromParts(Part.fromText("tool ended the invocation")))); + InvocationContext context = resumableContext(eventSession, "inv"); + + context.populateInvocationAgentStates(); + + assertThat(context.endOfAgents()).doesNotContainEntry("a", true); + } + + // An after-agent callback emits content after the end-of-agent marker; Python reopens the agent + // on it rather than leaving a resume with nothing to run. + @Test + public void populateInvocationAgentStates_contentAfterEndOfAgent_reopensAgent() { + Session eventSession = Session.builder("s").build(); + eventSession + .events() + .add(agentEvent("inv", "a", EventActions.builder().endOfAgent(true).build(), null)); + eventSession + .events() + .add(agentEvent("inv", "a", null, Content.fromParts(Part.fromText("callback output")))); + InvocationContext context = resumableContext(eventSession, "inv"); + + context.populateInvocationAgentStates(); + + assertThat(context.agentStates()).containsKey("a"); + assertThat(context.endOfAgents()).containsEntry("a", false); + } + + @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..d017860ce 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,298 @@ public void run_withExampleTool_doesNotAddFunctionDeclarations() { var config = request.config().get(); assertThat(config.tools().isPresent()).isFalse(); } + + // ---- Resumability: resume into a transferred sub-agent (ResumableLlmAgentTest parity). ---- + + 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(); + // Recorded in the invocation state too, so a later resume no-ops instead of re-running root. + assertThat(context.endOfAgents()).containsEntry("root", true); + } + + @Test + public void runAsync_resumeUserResponseWithNoMatchingCall_throws() { + // A user function response matching no prior call is rejected, not silently continued. + 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 sub-agent re-pauses, so root must not mark end-of-agent. + 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, not the action alone. + 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"); + } + + // A user event answering no call must not strand the transfer: the sub-agent still owns the turn. + @Test + public void runAsync_resumeAfterTransfer_userEventWithoutFunctionResponse_resumesSubAgent() { + LlmAgent sub = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("sub resumes"))) + .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")); + var unused = + context + .sessionService() + .appendEvent( + context.session(), + Event.builder() + .id("u1") + .invocationId("inv") + .author("user") + .content(Content.fromParts(Part.fromText("carry on"))) + .build()) + .blockingGet(); + context.setAgentState("root", ImmutableMap.of(), /* endOfAgent= */ false); + + List events = root.runAsync(context).toList().blockingGet(); + + assertThat(simplifyEvents(events)).contains("sub: sub resumes"); + assertThat(rootLlm.getRequests()).isEmpty(); + } + + // Mirrors AgentTransferTest.transferTargets_ordersSubAgentsThenParentThenPeers: the flow keeps a + // second copy of this rule that the agents package cannot reach, so both copies are pinned. + @Test + public void transferTargets_ordersSubAgentsThenParentThenPeers() { + LlmAgent child = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("x"))).name("child").build(); + LlmAgent peerA = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("x"))).name("peer_a").build(); + LlmAgent subject = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("x"))) + .name("subject") + .subAgents(child) + .build(); + LlmAgent peerB = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("x"))).name("peer_b").build(); + LlmAgent unused = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("x"))) + .name("parent") + .subAgents(peerA, subject, peerB) + .build(); + + assertThat(subject.transferTargets().stream().map(BaseAgent::name)) + .containsExactly("child", "parent", "peer_a", "peer_b") + .inOrder(); + } + + @Test + public void transferTargets_noLlmAgentParent_returnsOnlySubAgents() { + LlmAgent child = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("x"))).name("child").build(); + LlmAgent subject = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("x"))) + .name("subject") + .subAgents(child) + .build(); + + assertThat(subject.transferTargets().stream().map(BaseAgent::name)).containsExactly("child"); + } + + @Test + public void runAsync_resumeFromTransferToMissingAgent_throws() { + // The transferred-to agent is gone: a resumed transfer must fail, as a live one does. + 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..6afd913e7 100644 --- a/core/src/test/java/com/google/adk/agents/LoopAgentTest.java +++ b/core/src/test/java/com/google/adk/agents/LoopAgentTest.java @@ -19,16 +19,21 @@ 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; +import static org.junit.Assert.assertThrows; import com.google.adk.agents.Callbacks.BeforeAgentCallback; 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 +245,140 @@ 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(); + } + + // Restarting the count would loop past the cap, so an unreadable value fails the resume instead. + @Test + public void runAsync_resumingWithNonNumericTimesLooped_throws() { + 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); + context.setAgentState( + "loop", + ImmutableMap.of("current_sub_agent", "sub", "times_looped", "two"), + /* endOfAgent= */ false); + + var resumed = loopAgent.runAsync(context).toList(); + + assertThrows(IllegalStateException.class, resumed::blockingGet); + } + + @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() { + // Guards the cap comparison timesLooped >= maxIterations from below the cap. + 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 ea3e524fa..2989acbcf 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,19 @@ 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 static org.junit.Assert.assertThrows; +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,11 +39,14 @@ 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; import com.google.genai.types.Part; +import java.util.HashMap; import java.util.List; +import java.util.Map; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -341,4 +351,160 @@ private static InvocationContext contextResumingBranchedCall( context.branch(contextBranch); return context; } + + // ---- 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"); + } + + // A null name is not "finished": Python's model rejects it, so a resume must not silently skip + // every sub-agent. + @Test + public void runAsync_resumingFromNullSubAgentName_throws() { + LlmAgent agent1 = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("a1 done"))) + .name("agent1") + .build(); + SequentialAgent sequentialAgent = + SequentialAgent.builder().name("seq").subAgents(agent1).build(); + InvocationContext context = createResumableInvocationContext(sequentialAgent); + Map checkpoint = new HashMap<>(); + checkpoint.put("current_sub_agent", null); + context.setAgentState("seq", checkpoint, /* endOfAgent= */ false); + + var resumed = sequentialAgent.runAsync(context).toList(); + + assertThrows(IllegalStateException.class, resumed::blockingGet); + } + + @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 fast-forwards from history; the resumable flow uses its own checkpoints. + @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(); + + // A legacy paused session: agent2's long-running call, now answered, with no checkpoints. + 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/flows/llmflows/AgentTransferTest.java b/core/src/test/java/com/google/adk/flows/llmflows/AgentTransferTest.java index 79552520b..833dafd4c 100644 --- a/core/src/test/java/com/google/adk/flows/llmflows/AgentTransferTest.java +++ b/core/src/test/java/com/google/adk/flows/llmflows/AgentTransferTest.java @@ -20,9 +20,11 @@ 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.truth.Truth.assertThat; +import com.google.adk.agents.BaseAgent; import com.google.adk.agents.InvocationContext; import com.google.adk.agents.LiveRequest; import com.google.adk.agents.LiveRequestQueue; @@ -86,6 +88,78 @@ public Single> runAsync(Map args, ToolContex } } + // Order is load-bearing: this list is rendered into the model's transfer instruction. + @Test + public void transferTargets_ordersSubAgentsThenParentThenPeers() { + LlmAgent child = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("unused"))) + .name("child") + .build(); + LlmAgent peerA = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("unused"))) + .name("peer_a") + .build(); + LlmAgent subject = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("unused"))) + .name("subject") + .subAgents(child) + .build(); + LlmAgent peerB = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("unused"))) + .name("peer_b") + .build(); + LlmAgent unused = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("unused"))) + .name("parent") + .subAgents(peerA, subject, peerB) + .build(); + + assertThat(AgentTransfer.transferTargets(subject).stream().map(BaseAgent::name)) + .containsExactly("child", "parent", "peer_a", "peer_b") + .inOrder(); + } + + @Test + public void transferTargets_disallowedParentAndPeers_returnsOnlySubAgents() { + LlmAgent child = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("unused"))) + .name("child") + .build(); + LlmAgent peer = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("unused"))).name("peer").build(); + LlmAgent subject = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("unused"))) + .name("subject") + .subAgents(child) + .disallowTransferToParent(true) + .disallowTransferToPeers(true) + .build(); + LlmAgent unused = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("unused"))) + .name("parent") + .subAgents(subject, peer) + .build(); + + assertThat(AgentTransfer.transferTargets(subject).stream().map(BaseAgent::name)) + .containsExactly("child"); + } + + @Test + public void transferTargets_noLlmAgentParent_returnsOnlySubAgents() { + LlmAgent child = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("unused"))) + .name("child") + .build(); + LlmAgent subject = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("unused"))) + .name("subject") + .subAgents(child) + .build(); + + assertThat(AgentTransfer.transferTargets(subject).stream().map(BaseAgent::name)) + .containsExactly("child"); + } + @Test public void exitLoopTool_exitsLoop() { Content generatedContent = 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..c4aec95c0 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,13 @@ public void getAskUserConfirmationFunctionCalls_eventWithConfirmationFunctionCal assertThat(result).containsExactly(confirmationCall1, confirmationCall2); } + @Test + public void hasPendingLongRunningCall_singleEventWithFunctionResponse_returnsFalse() { + // One event cannot hold both a paused call and its response, so nothing is pending. + 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 +575,36 @@ public void hasPendingLongRunningCall_emptyList_returnsFalse() { assertThat(Functions.hasPendingLongRunningCall(ImmutableList.of())).isFalse(); } + @Test + public void hasPendingLongRunningCall_list_matchingResponse_stillReturnsTrue() { + // This legacy predicate does not match responses to calls; StepResume does that when resumable. + 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 index 7f0302758..b00583053 100644 --- a/core/src/test/java/com/google/adk/runner/RunnerLegacyResumabilityTest.java +++ b/core/src/test/java/com/google/adk/runner/RunnerLegacyResumabilityTest.java @@ -16,28 +16,36 @@ 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.apps.App; -import com.google.adk.apps.ResumabilityConfig; 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; @@ -45,9 +53,9 @@ import io.opentelemetry.api.trace.Tracer; import io.opentelemetry.sdk.testing.junit4.OpenTelemetryRule; import io.reactivex.rxjava3.core.Flowable; -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; @@ -62,6 +70,7 @@ * 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(); @@ -79,33 +88,396 @@ public void tearDown() { Tracing.setTracerForTesting(originalTracer); } - public static class Tools { - private Tools() {} + // Regression: after a transfer, a finished-check keyed on the root wedged later plain-text turns. + @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); - public static ImmutableMap echoTool(String message) { - return ImmutableMap.of("message", message); + // Turn 1 transfers; turns 2-5 must each be answered by the sub-agent, not 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); } + } - // A long-running tool awaiting an external result has nothing to return yet; FunctionTool - // coerces the absent return into an empty response. - @SuppressWarnings("unused") // Invoked reflectively by FunctionTool. - public static @Nullable ImmutableMap pendingTool(String message) { - return null; - } + // The shim must not re-attach to a checkpoint-less session: it starts a new invocation instead. + @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 checkpoints, as a pre-checkpoint session looks. + 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(); - static final AtomicInteger pendingProgressToolCalls = new AtomicInteger(0); + ImmutableList secondTurn = runTurn(runner, session, "second turn"); - // A long-running tool that reports progress: it returns a non-empty "pending" status on the - // initial call. Counts executions so a test can assert it runs exactly once across turns. - @SuppressWarnings("unused") // Invoked reflectively by FunctionTool. - public static ImmutableMap pendingProgressTool(String message) { - pendingProgressToolCalls.incrementAndGet(); - return ImmutableMap.of("status", "pending"); - } + // The agent runs under a fresh invocation id, not the checkpoint-less prior one. + assertThat(simplifyEvents(secondTurn)).contains("agent: second answer"); + assertThat( + secondTurn.stream() + .map(Event::invocationId) + .filter(priorInvocationId::equals) + .collect(toImmutableList())) + .isEmpty(); + } + + // The shim selects the legacy flow but does not resume on plain text; that starts a new one. + @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 continuation opens a new invocation and the agent plans again. + 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(); + } + + // The shim resumes only unfinished invocations, so text after a finished turn starts a new one. + @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); + } + + // A pause nested in a SequentialAgent is 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"); + + // 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 value-returning long-running tool still pauses BaseLlmFlow after one model call. + @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 legacy LoopAgent body stops on a pending 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: a mid-sequence pause restarts the sequence and persists no checkpoint. + @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: a grandchild-authored pause also restarts the sequence rather than resuming. + @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 shim, a plain-text continuation starts a new invocation and still merges stateDelta. + @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 - @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). public void runAsync_withToolConfirmation_inSequentialAgent_runsLaterSubAgentsAfterResume_legacyShim() { LlmAgent agentA = @@ -134,23 +506,10 @@ public static ImmutableMap pendingProgressTool(String message) { .name("workflow_agent") .subAgents(ImmutableList.of(agentA, agentB, agentC)) .build(); - Runner runner = - Runner.builder() - .app( - App.builder() - .name("test") - .rootAgent(workflowAgent) - .resumabilityConfig( - ResumabilityConfig.builder().plainTextContinuationAutoResume(true).build()) - .build()) - .build(); - Session session = runner.sessionService().createSession("test", "user").blockingGet(); + Runner runner = shimRunner(workflowAgent); + Session session = newSession(runner); - List eventsBeforeConfirmation = - runner - .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) - .toList() - .blockingGet(); + 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"); @@ -163,21 +522,22 @@ public static ImmutableMap pendingProgressTool(String message) { .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(); + 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)) @@ -189,7 +549,6 @@ public static ImmutableMap pendingProgressTool(String message) { } @Test - @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). public void runAsync_withLongRunningCall_inSequentialAgent_runsLaterSubAgentsAfterResume_legacyShim() { LlmAgent agentA = @@ -222,23 +581,10 @@ public static ImmutableMap pendingProgressTool(String message) { .name("workflow_agent") .subAgents(ImmutableList.of(agentA, agentB, agentC)) .build(); - Runner runner = - Runner.builder() - .app( - App.builder() - .name("test") - .rootAgent(workflowAgent) - .resumabilityConfig( - ResumabilityConfig.builder().plainTextContinuationAutoResume(true).build()) - .build()) - .build(); - Session session = runner.sessionService().createSession("test", "user").blockingGet(); + Runner runner = shimRunner(workflowAgent); + Session session = newSession(runner); - List eventsBeforeResume = - runner - .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) - .toList() - .blockingGet(); + 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 @@ -247,21 +593,8 @@ public static ImmutableMap pendingProgressTool(String message) { 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(); + 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)) @@ -270,7 +603,6 @@ public static ImmutableMap pendingProgressTool(String message) { } @Test - @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). public void runAsync_withLongRunningCall_resumable_pausesAfterSingleModelCall_legacyShim() { TestLlm testLlm = createTestLlm( @@ -290,23 +622,10 @@ public void runAsync_withLongRunningCall_resumable_pausesAfterSingleModelCall_le /* requireConfirmation= */ false, /* isLongRunning= */ true)) .build(); - Runner runner = - Runner.builder() - .app( - App.builder() - .name("test") - .rootAgent(agent) - .resumabilityConfig( - ResumabilityConfig.builder().plainTextContinuationAutoResume(true).build()) - .build()) - .build(); - Session session = runner.sessionService().createSession("test", "user").blockingGet(); + Runner runner = shimRunner(agent); + Session session = newSession(runner); - List events = - runner - .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) - .toList() - .blockingGet(); + 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); @@ -314,7 +633,6 @@ public void runAsync_withLongRunningCall_resumable_pausesAfterSingleModelCall_le } @Test - @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). public void runAsync_loopAgentWithLongRunningSubAgent_resumable_stopsAfterFirstIteration_legacyShim() { AtomicInteger calls = new AtomicInteger(); @@ -342,30 +660,16 @@ public void runAsync_withLongRunningCall_resumable_pausesAfterSingleModelCall_le .subAgents(ImmutableList.of(inner)) .maxIterations(3) .build(); - Runner runner = - Runner.builder() - .app( - App.builder() - .name("test") - .rootAgent(loop) - .resumabilityConfig( - ResumabilityConfig.builder().plainTextContinuationAutoResume(true).build()) - .build()) - .build(); - Session session = runner.sessionService().createSession("test", "user").blockingGet(); + Runner runner = shimRunner(loop); + Session session = newSession(runner); - List unused = - runner - .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) - .toList() - .blockingGet(); + ImmutableList unused = runTurn(runner, session, "from user"); // Paused after the first iteration: one model call, not maxIterations. assertThat(loopLlm.getRequests()).hasSize(1); } @Test - @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). public void runAsync_parallelAgentWithLongRunningBranch_resumable_otherBranchCompletes_legacyShim() { TestLlm longRunningLlm = @@ -392,23 +696,10 @@ public void runAsync_withLongRunningCall_resumable_pausesAfterSingleModelCall_le .name("parallel") .subAgents(ImmutableList.of(longRunningBranch, plainBranch)) .build(); - Runner runner = - Runner.builder() - .app( - App.builder() - .name("test") - .rootAgent(parallel) - .resumabilityConfig( - ResumabilityConfig.builder().plainTextContinuationAutoResume(true).build()) - .build()) - .build(); - Session session = runner.sessionService().createSession("test", "user").blockingGet(); + Runner runner = shimRunner(parallel); + Session session = newSession(runner); - List events = - runner - .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) - .toList() - .blockingGet(); + 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); 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..c612c52e5 --- /dev/null +++ b/core/src/test/java/com/google/adk/runner/RunnerResumabilityTest.java @@ -0,0 +1,1864 @@ +/* + * 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.reloadSession; +import static com.google.adk.testing.ResumabilityTestUtils.resumableRunner; +import static com.google.adk.testing.ResumabilityTestUtils.resume; +import static com.google.adk.testing.ResumabilityTestUtils.resumeFlowable; +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) { + // CALLS_REAL_METHODS keeps the default Maybe.empty() returns; a bare mock would return null. + 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); + } + + // A confirmation answer resumes B in SequentialAgent(A, B, C): C runs, A does not re-run. + @Test + public void runAsync_withToolConfirmation_inSequentialAgent_runsLaterSubAgentsAfterResume() { + LlmAgent agentA = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("agent A done"))) + .name("a_agent") + .build(); + // B pauses right after requesting confirmation, so one 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 runs the tool, then C runs; A is not re-run. + 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(); + } + + // A pending long-running call pauses SequentialAgent(A, B, C) at B; on resume C runs, not A. + @Test + public void runAsync_withLongRunningCall_inSequentialAgent_runsLaterSubAgentsAfterResume() { + LlmAgent agentA = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("agent A done"))) + .name("a_agent") + .build(); + // B pauses right after the no-result long-running call, so one 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. + 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. + 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 LoopAgent(w1, w2) paused on w1 resumes w1, then advances the loop to w2 and closes it. + @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, then the loop advances to w2 and closes; w1 is not re-run. + 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 transfer runs the sub-agent nested, as in Python: the root closes only after it completes. + @Test + public void runAsync_resumable_transferToSubAgent_closesRootAfterSubAgentCompletes() { + 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 sub-agent runs inside the root's turn; the root closes last. + 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={})", + "sub_agent_1: response1", + "sub_agent_1: end_of_agent", + "root_agent: 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 transferred-to sub-agent that pauses on a long-running call resumes in the same invocation. + @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(); + + // The sub-agent runs nested and pauses, so neither it nor the root closes. + ImmutableList pausedEvents = simplifyResumableEvents(pausedTurn); + assertThat(pausedEvents) + .contains("sub_agent_1: FunctionCall(name=pendingTool, args={message=hello})"); + assertThat(pausedEvents).doesNotContain("root_agent: end_of_agent"); + 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"); + // The root is an LlmAgent ancestor of the paused agent, so the resume re-enters it rather than + // routing past it, and it closes itself once the sub-agent finishes. + assertThat(simplifyResumableEvents(resumed)).contains("root_agent: end_of_agent"); + } + + // An invocation paused on two long-running calls stays paused until both are answered. + @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 calls are issued and the invocation pauses after one model call. + 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 long-running tool that returns a value resolves in-turn, so the flow summarizes, not 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"); + + // Summarized in-turn with no pause, and no end-of-agent: 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(); + } + + // A resume runs the whole plugin bracket, not only on-event, as the new-invocation path does. + @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"); + // Twice: once on the initial turn, once on the resume turn, which used to run only on-event. + 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 call keeps the invocation paused: plain text and no message are both inert. + @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 returns instead of 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 silent pause inside a nested Sequential leaves the outer checkpoint on the inner workflow. + @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 before advancing the outer one; nothing re-runs. + 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 sequence advances past it. + @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")); + + assertThat(simplifyEvents(resumed)).contains("leaf_agent: leaf resumed"); + // The enclosing sequence then advances past the parallel block. + assertThat(simplifyEvents(resumed)).contains("next_agent: next ran"); + } + + // A leaf under a ParallelAgent paused on two calls, answered once, stays paused. + @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(); + } + + // Two paused sibling leaves: answering one resumes it while the other holds the sequence. + @Test + public void runAsync_resume_parallelSiblingsPartiallyAnswered_answeredSiblingResumes() { + LlmAgent leafA = + createTestAgentBuilder( + createTestLlm( + createFunctionCallLlmResponse( + "c1", "pendingTool", ImmutableMap.of("message", "a")), + createTextLlmResponse("a resumed"))) + .name("leaf_a") + .tools(pendingFunctionTool()) + .build(); + LlmAgent leafB = + createTestAgentBuilder( + createTestLlm( + createFunctionCallLlmResponse( + "c2", "pendingTool", ImmutableMap.of("message", "b")), + createTextLlmResponse("b resumed"))) + .name("leaf_b") + .tools(pendingFunctionTool()) + .build(); + ParallelAgent parallel = + ParallelAgent.builder() + .name("parallel_agent") + .subAgents(ImmutableList.of(leafA, leafB)) + .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"); + + // Answer only leaf_a's call; leaf_b's stays unanswered. + ImmutableList resumed = + answerCall(runner, session, "c1", "pendingTool", ImmutableMap.of("message", "a")); + + assertThat(simplifyEvents(resumed)).contains("leaf_a: a resumed"); + assertThat(simplifyEvents(resumed)).doesNotContain("leaf_b: b resumed"); + // leaf_b is still pending, so the sequence must not advance past the parallel block. + assertThat(simplifyEvents(resumed)).doesNotContain("next_agent: next ran"); + // leaf_b must stay resumable, so nothing above it may close. + assertNoEndOfAgent(resumed, "leaf_b"); + assertNoEndOfAgent(resumed, "parallel_agent"); + assertNoEndOfAgent(resumed, "seq_agent"); + assertEndOfAgent(resumed, "leaf_a"); + } + + // With the shim off, a plain-text continuation after a pause starts a new invocation. + @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); + } + + // With resumability off, a completed LlmAgent emits no end-of-agent checkpoint. + @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 completed resumable LlmAgent emits a trailing end-of-agent; resuming past it is a no-op. + @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 end-of-agent checkpoint is how 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 matching invocation and merges the stateDelta. + @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 resume runs a persisted, unexecuted call, not the model. + @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 scripted model only returns text, so an echoTool response can only come from replay. + assertThat( + resumed.stream() + .flatMap(event -> event.functionResponses().stream()) + .map(response -> response.name().orElse("")) + .collect(toImmutableList())) + .contains("echoTool"); + // Replay precedes the model, so the first request already carries the tool's response. + assertThat(testLlm.getRequests().get(0).contents().toString()).contains("echoTool"); + } + + // Python restricts a transfer to the agent's declared targets; a replay must not get around it. + @Test + public void resume_replayedTransferToUndeclaredAgent_fails() { + TestLlm peerLlm = createTestLlm(createTextLlmResponse("peer child should not run")); + LlmAgent peerChild = createTestAgentBuilder(peerLlm).name("peer_child").build(); + LlmAgent peer = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("peer"))) + .name("peer") + .subAgents(peerChild) + .build(); + // No sub-agents and no peer transfers, so peer_child is not a declared target of agent. + LlmAgent agent = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("should not be reached"))) + .name("agent") + .disallowTransferToPeers(true) + .build(); + LlmAgent root = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("root"))) + .name("root") + .subAgents(agent, peer) + .build(); + Runner runner = resumableRunner(root); + Session session = newSession(runner); + String invocationId = "inv-bad-transfer"; + var unusedUser = + runner + .sessionService() + .appendEvent( + session, + Event.builder() + .id("u1") + .invocationId(invocationId) + .author("user") + .content(Content.fromParts(Part.fromText("go"))) + .build()) + .blockingGet(); + // A transfer call persisted but never executed, naming an agent this one may not transfer to. + 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("transfer_to_agent") + .args(ImmutableMap.of("agent_name", "peer_child")) + .build()) + .build())) + .build()) + .blockingGet(); + + resumeFlowable(runner, session, invocationId, /* newMessage= */ null, /* stateDelta= */ null) + .test() + .assertError( + error -> + error instanceof IllegalStateException + && error.getMessage().contains("is not allowed to transfer")); + assertThat(peerLlm.getRequests()).isEmpty(); + } + + // Python parity: a stateDelta resumed with no message is persisted as a content-less event. + @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: a function response resumes an older invocation, not the latest. + @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 state from history. + @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")); + } + + // Python parity: resume text is dropped when the invocation already has a user message. + @Test + public void resume_withTextMessage_invocationAlreadyHasUserMessage_doesNotAppendIt() { + 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(); + // Only the original user message remains: the resume text was not appended. + assertThat(reloaded.events()).hasSize(1); + assertThat( + Iterables.getLast(reloaded.events()) + .content() + .flatMap(Content::parts) + .get() + .get(0) + .text()) + .hasValue("hi"); + } + + // The same resume does append the text when the invocation has no user message to duplicate. + @Test + public void resume_withTextMessage_invocationHasNoUserMessage_appendsIt() { + TestBaseAgent agent = + new TestBaseAgent("test_agent", "desc", () -> Flowable.empty(), null, null, null); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + Object unusedSeed = + runner + .sessionService() + .appendEvent( + session, + Event.builder() + .id("a1") + .invocationId("test-inv") + .author("test_agent") + .content(Content.fromParts(Part.fromText("agent spoke first"))) + .build()) + .blockingGet(); + + Object unused = + resume( + runner, + session, + "test-inv", + Content.fromParts(Part.fromText("New message")), + /* stateDelta= */ null); + + Session reloaded = reloadSession(runner, session); + 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)); + } + + // With no invocation id, an orphan function response is rejected rather than fed to the model. + @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); + } + + // Python defaults a role-less message to "user"; without it the LLM request drops the response. + @Test + public void resume_messageWithoutRole_persistsItAsAUserTurnTheModelCanSee() { + TestLlm testLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hi")), + createTextLlmResponse("resumed answer")); + LlmAgent agent = + createTestAgentBuilder(testLlm).name("agent").tools(pendingFunctionTool()).build(); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + var unused = runTurn(runner, session, "hi"); + + // Built with no role, the shape a caller assembling a function response naturally produces. + Content roleless = + Content.builder() + .parts( + ImmutableList.of( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("lro_call_id") + .name("pendingTool") + .response(ImmutableMap.of("status", "done"))) + .build())) + .build(); + var resumed = + resume(runner, session, /* invocationId= */ null, roleless, /* stateDelta= */ null); + + assertThat(simplifyEvents(resumed)).contains("agent: resumed answer"); + Content lastRequestContent = + Iterables.getLast(Iterables.getLast(testLlm.getRequests()).contents()); + assertThat(lastRequestContent.role()).hasValue("user"); + assertThat(lastRequestContent.parts().get().get(0).functionResponse()).isPresent(); + } + + // The ordinary runAsync overload routes into the resumable flow too, so it needs the same + // role default as the resume overload. + @Test + public void runAsync_resumableApp_messageWithoutRole_reachesTheModel() { + TestLlm testLlm = createTestLlm(createTextLlmResponse("answered")); + LlmAgent agent = createTestAgentBuilder(testLlm).name("agent").build(); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + Content roleless = + Content.builder().parts(ImmutableList.of(Part.fromText("no role here"))).build(); + + var unused = + runner + .runAsync("user", session.id(), roleless, RunConfig.builder().build()) + .toList() + .blockingGet(); + + Content sent = Iterables.getLast(Iterables.getLast(testLlm.getRequests()).contents()); + assertThat(sent.role()).hasValue("user"); + assertThat(sent.parts().get().get(0).text()).hasValue("no role here"); + } + + // The declared target wins over a same-named agent elsewhere in the tree, as in Python. + @Test + public void resume_replayedTransfer_prefersDeclaredTargetOverSameNamedAgent() { + TestLlm declaredLlm = createTestLlm(createTextLlmResponse("declared ran")); + TestLlm strangerLlm = createTestLlm(createTextLlmResponse("stranger ran")); + // Two agents share the name "helper": one is agent's own sub-agent, one is under a peer. + LlmAgent declaredHelper = createTestAgentBuilder(declaredLlm).name("helper").build(); + LlmAgent agent = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("unused"))) + .name("agent") + .subAgents(declaredHelper) + .build(); + LlmAgent stranger = createTestAgentBuilder(strangerLlm).name("helper").build(); + LlmAgent peer = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("unused"))) + .name("peer") + .subAgents(stranger) + .build(); + LlmAgent root = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("unused"))) + .name("root") + // peer first, so a tree-wide search reaches its "helper" before the declared one. + .subAgents(peer, agent) + .build(); + Runner runner = resumableRunner(root); + Session session = newSession(runner); + String invocationId = "inv-dup-name"; + var unusedUser = + runner + .sessionService() + .appendEvent( + session, + Event.builder() + .id("u1") + .invocationId(invocationId) + .author("user") + .content(Content.fromParts(Part.fromText("go"))) + .build()) + .blockingGet(); + 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("transfer_to_agent") + .args(ImmutableMap.of("agent_name", "helper")) + .build()) + .build())) + .build()) + .blockingGet(); + + var unused = + resume(runner, session, invocationId, /* newMessage= */ null, /* stateDelta= */ null); + + assertThat(declaredLlm.getRequests()).isNotEmpty(); + assertThat(strangerLlm.getRequests()).isEmpty(); + } + + // A message mixing text with a function response resumes on the response; the text is ignored. + @Test + public void resume_functionResponseMixedWithText_resumesIgnoringTheText() { + 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")); + + ImmutableList resumed = + ImmutableList.copyOf( + runner + .runAsync( + "user", + session.id(), + /* invocationId= */ null, + mixed, + RunConfig.builder().build(), + /* stateDelta= */ null) + .blockingIterable()); + + assertThat(simplifyEvents(resumed)).contains("agent: done"); + } + + // Responses answering calls from two invocations are rejected, not resolved to the newest. + // An id-less function response cannot be routed to a call, so the resume is rejected rather than + // opening a new invocation carrying an orphan response, as in Python. + @Test + public void resume_functionResponseWithoutId_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 idLess = + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .name("pendingTool") + .response(ImmutableMap.of("message", "hi"))) + .build()); + + runner + .runAsync( + "user", + session.id(), + /* invocationId= */ null, + idLess, + RunConfig.builder().build(), + /* stateDelta= */ null) + .test() + .assertError(IllegalArgumentException.class); + } + + @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 recovers the invocation's original user message. + @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 under a SequentialAgent paused on two calls, answered once, must not look completed. + @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"); + } + + // Python parity: each ParallelAgent sub-agent runs on its own parallel. branch. + @Test + public void runAsync_resumable_parallelSubAgents_runOnTheirOwnBranches() { + LlmAgent branchA = + createTestAgentBuilder(pauseThenSay("call_a", "branch A resumed")) + .name("branch_a") + .tools(pendingFunctionTool()) + .build(); + LlmAgent branchB = textAgent("branch_b", "branch B done"); + ParallelAgent parallel = + ParallelAgent.builder() + .name("parallel_agent") + .subAgents(ImmutableList.of(branchA, branchB)) + .build(); + Runner runner = resumableRunner(parallel); + Session session = newSession(runner); + + runTurn(runner, session, "start"); + Session reloaded = reloadSession(runner, session); + + assertThat(branchesOf(reloaded, "branch_a")).containsExactly("parallel_agent.branch_a"); + assertThat(branchesOf(reloaded, "branch_b")).containsExactly("parallel_agent.branch_b"); + // The answer inherits the branch of the call it answers, so it lands on branch A alone. + answerCall(runner, session, "call_a", "pendingTool", ImmutableMap.of("message", "hello")); + assertThat(branchesOf(reloadSession(runner, session), "user")) + .contains("parallel_agent.branch_a"); + } + + /** The distinct branch strings carried by {@code author}'s events in {@code session}. */ + private static ImmutableList branchesOf(Session session, String author) { + return session.events().stream() + .filter(event -> author.equals(event.author())) + .map(event -> event.branch().orElse("-")) + .distinct() + .collect(toImmutableList()); + } + + // An answer on a parallel sub-branch the LlmAgent cannot see must still resume the leaf. + @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"); + } + + // A transfer inside a workflow: the sequence must advance to the sub-agent after the transferring + // agent once the transferred-to leaf resumes, instead of stranding on the leaf. + @Test + public void runAsync_resume_transferInsideSequential_sequenceAdvancesAfterLeafResumes() { + LlmAgent leaf = + createTestAgentBuilder( + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hi")), + createTextLlmResponse("leaf resumed"))) + .name("leaf_agent") + .tools(pendingFunctionTool()) + .build(); + LlmAgent transferring = + createTestAgentBuilder( + createTestLlm( + createFunctionCallLlmResponse( + "transfer_call", + "transfer_to_agent", + ImmutableMap.of("agent_name", "leaf_agent")))) + .name("transferring_agent") + .subAgents(leaf) + .build(); + LlmAgent next = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("next ran"))) + .name("next_agent") + .build(); + SequentialAgent root = + SequentialAgent.builder() + .name("seq_agent") + .subAgents(ImmutableList.of(transferring, 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")); + + assertThat(simplifyEvents(resumed)).contains("leaf_agent: leaf resumed"); + // The enclosing sequence must not stall on the leaf. + assertThat(simplifyEvents(resumed)).contains("next_agent: next ran"); + } + + /** 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 message throws instead of calling the model. + @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); + } + + // An explicit invocation id skips message resolution, so the orphan check must fire here too. + @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 call's branch. + @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 stops a resumable LoopAgent after the iteration, as 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); + } + + // A pending long-running call pauses only its own ParallelAgent branch, as Python ADK v1 does. + @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/sessions/InMemorySessionServiceTest.java b/core/src/test/java/com/google/adk/sessions/InMemorySessionServiceTest.java index 58c445641..2f3a5bcae 100644 --- a/core/src/test/java/com/google/adk/sessions/InMemorySessionServiceTest.java +++ b/core/src/test/java/com/google/adk/sessions/InMemorySessionServiceTest.java @@ -16,13 +16,16 @@ package com.google.adk.sessions; import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; import com.google.adk.events.Event; import com.google.adk.events.EventActions; +import com.google.adk.sessions.InMemorySessionService.DuplicateSessionIdBehavior; import io.reactivex.rxjava3.core.Single; import java.lang.reflect.Field; import java.time.Instant; import java.util.HashMap; +import java.util.Map; import java.util.Optional; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; @@ -340,4 +343,88 @@ public void getSession_numRecentEventsAndAfterTimestamp_appliesBothFilters() { .containsExactly(300L, 400L, 500L) .inOrder(); } + + @Test + public void createSession_duplicateSessionId_byDefaultReplacesExistingSession() { + InMemorySessionService sessionService = new InMemorySessionService(); + Map state = new HashMap<>(); + state.put("sessionKey", "sessionValue"); + Session session = sessionService.createSession("app", "user", state, "session1").blockingGet(); + var unused = + sessionService.appendEvent(session, Event.builder().timestamp(100).build()).blockingGet(); + + Session replacement = + sessionService.createSession("app", "user", new HashMap<>(), "session1").blockingGet(); + + assertThat(replacement.events()).isEmpty(); + assertThat(replacement.state()).doesNotContainKey("sessionKey"); + Session retrieved = + sessionService.getSession("app", "user", "session1", Optional.empty()).blockingGet(); + assertThat(retrieved.events()).isEmpty(); + assertThat(retrieved.state()).doesNotContainKey("sessionKey"); + } + + @Test + public void createSession_duplicateSessionId_whenRejecting_errorsAndKeepsExistingSession() { + InMemorySessionService sessionService = + new InMemorySessionService(DuplicateSessionIdBehavior.REJECT); + Map state = new HashMap<>(); + state.put("sessionKey", "sessionValue"); + Session session = sessionService.createSession("app", "user", state, "session1").blockingGet(); + var unused = + sessionService.appendEvent(session, Event.builder().timestamp(100).build()).blockingGet(); + + Single duplicate = + sessionService.createSession("app", "user", new HashMap<>(), "session1"); + + SessionException exception = assertThrows(SessionException.class, duplicate::blockingGet); + assertThat(exception).hasMessageThat().isEqualTo(SessionException.SESSION_ALREADY_EXISTS); + Session retrieved = + sessionService.getSession("app", "user", "session1", Optional.empty()).blockingGet(); + assertThat(retrieved.events().stream().map(Event::timestamp)).containsExactly(100L); + assertThat(retrieved.state()).containsEntry("sessionKey", "sessionValue"); + } + + @Test + public void createSession_distinctSessionIds_whenRejecting_createsBoth() { + InMemorySessionService sessionService = + new InMemorySessionService(DuplicateSessionIdBehavior.REJECT); + + Session first = + sessionService.createSession("app", "user", new HashMap<>(), "session1").blockingGet(); + Session second = + sessionService.createSession("app", "user", new HashMap<>(), "session2").blockingGet(); + + assertThat(first.id()).isEqualTo("session1"); + assertThat(second.id()).isEqualTo("session2"); + assertThat(sessionService.listSessions("app", "user").blockingGet().sessions()).hasSize(2); + } + + @Test + public void createSession_afterDelete_whenRejecting_reusesTheSameId() { + InMemorySessionService sessionService = + new InMemorySessionService(DuplicateSessionIdBehavior.REJECT); + var unused = + sessionService.createSession("app", "user", new HashMap<>(), "session1").blockingGet(); + sessionService.deleteSession("app", "user", "session1").blockingAwait(); + + Session recreated = + sessionService.createSession("app", "user", new HashMap<>(), "session1").blockingGet(); + + assertThat(recreated.id()).isEqualTo("session1"); + } + + @Test + public void createSession_sameSessionIdDifferentUser_whenRejecting_isAllowed() { + InMemorySessionService sessionService = + new InMemorySessionService(DuplicateSessionIdBehavior.REJECT); + var unused = + sessionService.createSession("app", "user-a", new HashMap<>(), "session1").blockingGet(); + + Session other = + sessionService.createSession("app", "user-b", new HashMap<>(), "session1").blockingGet(); + + assertThat(other.id()).isEqualTo("session1"); + assertThat(other.userId()).isEqualTo("user-b"); + } } 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..d52ef93da 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,107 @@ 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); + } + + // rawEvent supplies only the keys "actions" lacks: a rawEvent carrying endOfAgent alone must not + // shadow an agentState that actions does carry. + @Test + public void fromApiEvent_rawEventMissingAgentState_doesNotShadowActions() { + 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")); + apiEvent.put("actions", actions); + apiEvent.put("rawEvent", ImmutableMap.of("actions", ImmutableMap.of("endOfAgent", true))); + + Event event = SessionJsonConverter.fromApiEvent(apiEvent); + + assertThat(event.actions().agentState()).isPresent(); + assertThat(event.actions().agentState().get()).containsEntry("current_sub_agent", "b_agent"); + assertThat(event.actions().endOfAgent()).isTrue(); + } + + // Python's Vertex session service keeps agentState and endOfAgent only under rawEvent, so a + // session it wrote must still load with its checkpoints. + @Test + public void fromApiEvent_agentStateOnlyUnderRawEvent_readsCheckpoint() { + 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"); + // The six-key actions block Python writes, with no agentState and no endOfAgent. + apiEvent.put("actions", new HashMap<>(ImmutableMap.of("stateDelta", new HashMap<>()))); + Map rawActions = new HashMap<>(); + rawActions.put("agentState", ImmutableMap.of("current_sub_agent", "b_agent")); + rawActions.put("endOfAgent", true); + apiEvent.put("rawEvent", ImmutableMap.of("actions", rawActions)); + + Event event = SessionJsonConverter.fromApiEvent(apiEvent); + + assertThat(event.actions().agentState()).isPresent(); + assertThat(event.actions().agentState().get()).containsEntry("current_sub_agent", "b_agent"); + assertThat(event.actions().endOfAgent()).isTrue(); + } + + @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() { + // Another runtime may write a non-map agentState; dropping it keeps the session loadable. + 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 diff --git a/core/src/test/java/com/google/adk/tools/BaseToolTest.java b/core/src/test/java/com/google/adk/tools/BaseToolTest.java index 960e8aab3..e78750a43 100644 --- a/core/src/test/java/com/google/adk/tools/BaseToolTest.java +++ b/core/src/test/java/com/google/adk/tools/BaseToolTest.java @@ -183,7 +183,7 @@ public void processLlmRequestWithGoogleSearchToolAddsToolToConfig() { } @Test - public void processLlmRequestWithLatestAliasAddsToolToConfig() { + public void processLlmRequestWithGoogleSearchToolAndLatestAliasAddsToolToConfig() { final GoogleSearchTool googleSearchTool = new GoogleSearchTool(); LlmRequest.Builder builder = LlmRequest.builder().model("gemini-flash-latest").build().toBuilder(); @@ -194,19 +194,23 @@ public void processLlmRequestWithLatestAliasAddsToolToConfig() { } @Test - public void processLlmRequestWithUnsupportedModelReturnsError() { - final GoogleSearchTool googleSearchTool = new GoogleSearchTool(); - LlmRequest.Builder builder = LlmRequest.builder().model("text-bison-001").build().toBuilder(); - Completable result = googleSearchTool.processLlmRequest(builder, null); - result.test().assertError(IllegalArgumentException.class); + public void processLlmRequestWithGoogleSearchToolAndNonGeminiModelAddsToolToConfig() { + GoogleSearchTool googleSearchTool = new GoogleSearchTool(); + LlmRequest.Builder builder = LlmRequest.builder().model("text-bison-001"); + Completable result = googleSearchTool.processLlmRequest(builder, /* toolContext= */ null); + result.test().assertComplete(); + assertThat(builder.build().config().get().tools().get()) + .containsExactly(Tool.builder().googleSearch(GoogleSearch.builder().build()).build()); } @Test - public void processLlmRequest_WithNullModel_ReturnsError() { - final GoogleSearchTool googleSearchTool = new GoogleSearchTool(); - LlmRequest.Builder builder = LlmRequest.builder().build().toBuilder(); - Completable result = googleSearchTool.processLlmRequest(builder, null); - result.test().assertError(IllegalArgumentException.class); + public void processLlmRequestWithGoogleSearchToolAndNoModelAddsToolToConfig() { + GoogleSearchTool googleSearchTool = new GoogleSearchTool(); + LlmRequest.Builder builder = LlmRequest.builder(); + Completable result = googleSearchTool.processLlmRequest(builder, /* toolContext= */ null); + result.test().assertComplete(); + assertThat(builder.build().config().get().tools().get()) + .containsExactly(Tool.builder().googleSearch(GoogleSearch.builder().build()).build()); } @Test @@ -238,6 +242,48 @@ public void processLlmRequestWithUrlContextToolAddsToolToConfig() { Tool.builder().urlContext(UrlContext.builder().build()).build()); } + @Test + public void processLlmRequestWithUrlContextToolAndLatestAliasAddsToolToConfig() { + UrlContextTool urlContextTool = new UrlContextTool(); + LlmRequest.Builder builder = LlmRequest.builder().model("gemini-flash-latest"); + Completable result = urlContextTool.processLlmRequest(builder, /* toolContext= */ null); + result.test().assertComplete(); + assertThat(builder.build().config().get().tools().get()) + .containsExactly(Tool.builder().urlContext(UrlContext.builder().build()).build()); + } + + @Test + public void processLlmRequestWithUrlContextToolAndResourcePathModelAddsToolToConfig() { + UrlContextTool urlContextTool = new UrlContextTool(); + LlmRequest.Builder builder = + LlmRequest.builder() + .model("projects/p/locations/us-central1/publishers/google/models/gemini-2.5-flash"); + Completable result = urlContextTool.processLlmRequest(builder, /* toolContext= */ null); + result.test().assertComplete(); + assertThat(builder.build().config().get().tools().get()) + .containsExactly(Tool.builder().urlContext(UrlContext.builder().build()).build()); + } + + @Test + public void processLlmRequestWithUrlContextToolAndNonGeminiModelAddsToolToConfig() { + UrlContextTool urlContextTool = new UrlContextTool(); + LlmRequest.Builder builder = LlmRequest.builder().model("text-bison-001"); + Completable result = urlContextTool.processLlmRequest(builder, /* toolContext= */ null); + result.test().assertComplete(); + assertThat(builder.build().config().get().tools().get()) + .containsExactly(Tool.builder().urlContext(UrlContext.builder().build()).build()); + } + + @Test + public void processLlmRequestWithUrlContextToolAndNoModelAddsToolToConfig() { + UrlContextTool urlContextTool = new UrlContextTool(); + LlmRequest.Builder builder = LlmRequest.builder(); + Completable result = urlContextTool.processLlmRequest(builder, /* toolContext= */ null); + result.test().assertComplete(); + assertThat(builder.build().config().get().tools().get()) + .containsExactly(Tool.builder().urlContext(UrlContext.builder().build()).build()); + } + private static InvocationContext.Builder testInvocationContext() { InvocationContext.Builder builder = InvocationContext.builder(); builder.agent(testAgent().build()); @@ -399,16 +445,6 @@ public void runAsync_withClassAndObjectMapper_convertsArguments() throws Excepti testObserver.assertValue(expected); } - @Test - public void testProcessLlmRequest_WithNoModel_DoesNotThrowsException() { - GoogleSearchTool tool = GoogleSearchTool.INSTANCE; - LlmRequest.Builder requestBuilder = LlmRequest.builder(); - - tool.processLlmRequest(requestBuilder, null); - - assertNotNull(requestBuilder); - } - public record TestToolArgs(int i, String s) {} @Test diff --git a/dev/README.md b/dev/README.md index 3f4e9a3d3..8ab2621f9 100644 --- a/dev/README.md +++ b/dev/README.md @@ -6,3 +6,34 @@ The UI and its assets are served under `/dev-ui/`, and both `/` and `/dev-ui` redirect there, keeping the query string. The assets are not served from the origin root: `/adk_favicon.svg` and the like return 404, and only the `/dev-ui/` form resolves. + +## Behind a reverse proxy + +If the server runs behind a reverse proxy that strips a path prefix (for +example, forwarding `https://gateway.example.com/my-app/` to `/`), set +`adk.web.backend-url` to the external URL that browsers use to reach the server: + +```properties +adk.web.backend-url=https://gateway.example.com/my-app +``` + +Setting `adk.web.backend-url`: + +- Sets `backendUrl` in `/dev-ui/assets/config/runtime-config.json` so the UI + sends API and WebSocket requests through the proxy URL. +- Prepends the URL's path (`/my-app`) to the `/` and `/dev-ui` entry redirects + (`/my-app/dev-ui/`). + +The value must be an absolute `http://` or `https://` URL without credentials, +query parameters, or a fragment. The UI parses values without an `http://` or +`https://` scheme as a WebSocket host, so a relative path like `/my-app` cannot +be used. If the application also sets `server.servlet.context-path`, include it +in `adk.web.backend-url`; when `adk.web.backend-url` has a path prefix, the +redirect uses that path directly instead of prepending the context path or +`X-Forwarded-Prefix`. + +When `adk.web.backend-url` is unset, `/` and `/dev-ui` redirect to `/dev-ui/` +and the bundled `backendUrl` value is preserved. If you instead enable +`server.forward-headers-strategy=framework`, make sure your reverse proxy strips +or overwrites incoming `Forwarded` and `X-Forwarded-*` headers from untrusted +clients. diff --git a/dev/src/main/java/com/google/adk/web/AdkWebServer.java b/dev/src/main/java/com/google/adk/web/AdkWebServer.java index 0a6035923..b4a4a92b8 100644 --- a/dev/src/main/java/com/google/adk/web/AdkWebServer.java +++ b/dev/src/main/java/com/google/adk/web/AdkWebServer.java @@ -25,6 +25,7 @@ import com.google.adk.memory.InMemoryMemoryService; import com.google.adk.sessions.BaseSessionService; import com.google.adk.sessions.InMemorySessionService; +import com.google.adk.web.config.BackendUrl; import com.google.adk.web.config.DevUiAssets; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -52,6 +53,15 @@ public class AdkWebServer implements WebMvcConfigurer { @Value("${adk.web.ui.dir:#{null}}") private String webUiDir; + @Value("${adk.web.backend-url:}") + private String backendUrlProperty; + + /** One bean, so this and the runtime-config endpoint cannot read the property differently. */ + @Bean + public BackendUrl backendUrl() { + return BackendUrl.from(backendUrlProperty); + } + @Bean public BaseSessionService sessionService() { // TODO: Add logic to select service based on config (e.g., DB URL) @@ -125,17 +135,25 @@ public void addResourceHandlers(ResourceHandlerRegistry registry) { } /** - * Configures simple automated controllers: "/" and "/dev-ui" both redirect to "/dev-ui/", which - * forwards to the UI's index.html. The trailing slash is required: index.html declares a {@code - * }, so served from "/dev-ui" the app resolves its own router path to "dev-ui" - * and matches none of its routes. The query string is carried across because the UI selects its - * agent from {@code ?app=} and the sample READMEs send users to the slashless "/dev-ui", so a - * redirect that dropped it would silently ignore the selection. + * Configures simple automated controllers: "/" and "/dev-ui" both redirect to the UI, at {@code + * adk.web.backend-url}'s path when that is set, and it forwards to index.html. The trailing slash + * is required: index.html declares a {@code }, so served from "/dev-ui" the app + * resolves its own router path to "dev-ui" and matches none of its routes. The query string is + * carried across because the UI selects its agent from {@code ?app=}. */ @Override public void addViewControllers(ViewControllerRegistry registry) { - registry.addRedirectViewController("/", "/dev-ui/").setKeepQueryParams(true); - registry.addRedirectViewController("/dev-ui", "/dev-ui/").setKeepQueryParams(true); + String prefix = backendUrl().pathPrefix(); + // The configured value is the public base, so do not stack the context path on it. + boolean contextRelative = prefix.isEmpty(); + registry + .addRedirectViewController("/", prefix + "/dev-ui/") + .setKeepQueryParams(true) + .setContextRelative(contextRelative); + registry + .addRedirectViewController("/dev-ui", prefix + "/dev-ui/") + .setKeepQueryParams(true) + .setContextRelative(contextRelative); registry.addViewController("/dev-ui/").setViewName("forward:/dev-ui/index.html"); } diff --git a/dev/src/main/java/com/google/adk/web/config/BackendUrl.java b/dev/src/main/java/com/google/adk/web/config/BackendUrl.java new file mode 100644 index 000000000..07f9527e7 --- /dev/null +++ b/dev/src/main/java/com/google/adk/web/config/BackendUrl.java @@ -0,0 +1,124 @@ +/* + * 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.web.config; + +import com.google.common.base.CharMatcher; +import com.google.common.base.Strings; +import java.net.URI; +import java.net.URISyntaxException; +import org.jspecify.annotations.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Parsed representation of {@code adk.web.backend-url}, shared by the runtime-config endpoint and + * the dev UI entry redirects. + */ +public final class BackendUrl { + + private static final Logger log = LoggerFactory.getLogger(BackendUrl.class); + + private static final BackendUrl UNSET = new BackendUrl("", ""); + + private final String value; + private final String pathPrefix; + + private BackendUrl(String value, String pathPrefix) { + this.value = value; + this.pathPrefix = pathPrefix; + } + + /** + * Parses {@code configured}. If the URL cannot be used by the dev UI, logs a warning, retains the + * trimmed value for {@code runtime-config.json}, and sets {@link #pathPrefix()} to empty. + */ + public static BackendUrl from(@Nullable String configured) { + String trimmed = Strings.nullToEmpty(configured).trim(); + if (trimmed.isEmpty()) { + return UNSET; + } + // Strip trailing slashes because the UI appends paths that already start with "/". + String normalized = CharMatcher.is('/').trimTrailingFrom(trimmed); + URI uri = parse(normalized); + String fault = faultIn(uri); + if (fault == null) { + return new BackendUrl(normalized, pathPrefixOf(uri)); + } + log.warn( + "adk.web.backend-url \"{}\" is not usable ({}), so the dev UI's redirect carries no" + + " prefix.", + withoutCredentials(trimmed), + fault); + return new BackendUrl(trimmed, ""); + } + + /** Returns the URL string to report in {@code runtime-config.json}, or empty when unset. */ + public String value() { + return value; + } + + /** + * Returns the URL path prefix to prepend to the entry redirect, or empty if unset or unusable. + * Preserves raw percent-encoding for use in a {@code Location} header. + */ + public String pathPrefix() { + return pathPrefix; + } + + private static @Nullable URI parse(String url) { + try { + return new URI(url); + } catch (URISyntaxException e) { + return null; + } + } + + private static @Nullable String faultIn(@Nullable URI uri) { + if (uri == null) { + return "it is not a valid URI"; + } + // The UI strips "http://" or "https://" case-sensitively when building its WebSocket URL. + if (!"http".equals(uri.getScheme()) && !"https".equals(uri.getScheme())) { + return "the scheme must be a lower-case http or https"; + } + if (uri.getHost() == null) { + return "it names no host"; + } + if (uri.getRawUserInfo() != null) { + return "it carries credentials, which would be served to every client"; + } + // The UI appends request paths directly onto backendUrl, so a query or fragment would corrupt + // the resulting URL. + if (uri.getRawQuery() != null || uri.getRawFragment() != null) { + return "it carries a query or fragment"; + } + return null; + } + + private static String pathPrefixOf(URI uri) { + String raw = uri.getRawPath(); + if (raw == null || raw.isEmpty()) { + return ""; + } + // Collapse leading slashes so "//segment" cannot be interpreted as a protocol-relative host. + return CharMatcher.is('/').trimTrailingFrom(raw.replaceAll("^/+", "/")); + } + + private static String withoutCredentials(String value) { + return value.replaceFirst("(?<=//)[^/@]*@", "***@"); + } +} diff --git a/dev/src/main/java/com/google/adk/web/config/DevUiAssets.java b/dev/src/main/java/com/google/adk/web/config/DevUiAssets.java index 48c742c8b..2ef291c48 100644 --- a/dev/src/main/java/com/google/adk/web/config/DevUiAssets.java +++ b/dev/src/main/java/com/google/adk/web/config/DevUiAssets.java @@ -20,11 +20,15 @@ import org.springframework.core.io.ResourceLoader; /** - * Where the dev UI's static assets live. Normalizes {@code adk.web.ui.dir} into a resource - * location, so callers that need it do not each do it differently. + * Where the dev UI's static assets live. Shared so the resource handler and the runtime-config + * endpoint resolve the same location; normalizing {@code adk.web.ui.dir} separately in each would + * diverge silently. */ public final class DevUiAssets { + /** The runtime config, relative to the asset root. */ + public static final String RUNTIME_CONFIG_PATH = "assets/config/runtime-config.json"; + private static final String CLASSPATH_ROOT = ResourceLoader.CLASSPATH_URL_PREFIX + "/browser/"; /** @@ -42,5 +46,10 @@ public static String assetRoot(@Nullable String webUiDir) { return location.endsWith("/") ? location : location + "/"; } + /** The location of a single asset, given relative to the asset root. */ + public static String assetLocation(@Nullable String webUiDir, String relativePath) { + return assetRoot(webUiDir) + relativePath; + } + private DevUiAssets() {} } diff --git a/dev/src/main/java/com/google/adk/web/controller/RuntimeConfigController.java b/dev/src/main/java/com/google/adk/web/controller/RuntimeConfigController.java new file mode 100644 index 000000000..7f538e2b6 --- /dev/null +++ b/dev/src/main/java/com/google/adk/web/controller/RuntimeConfigController.java @@ -0,0 +1,104 @@ +/* + * 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.web.controller; + +import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.adk.web.config.BackendUrl; +import com.google.adk.web.config.DevUiAssets; +import java.io.IOException; +import java.io.InputStream; +import java.util.LinkedHashMap; +import java.util.Map; +import org.jspecify.annotations.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.beans.factory.annotation.Value; +import org.springframework.core.io.Resource; +import org.springframework.core.io.ResourceLoader; +import org.springframework.http.CacheControl; +import org.springframework.http.MediaType; +import org.springframework.http.ResponseEntity; +import org.springframework.web.bind.annotation.GetMapping; +import org.springframework.web.bind.annotation.RestController; + +/** + * Serves the dev UI's runtime configuration, shadowing the copy bundled in the static assets so + * {@code adk.web.backend-url} can point the UI at the address browsers reach this server on. The + * bundled document is merged rather than replaced, so keys the UI gains in a later bundle survive. + */ +@RestController +public class RuntimeConfigController { + + private static final Logger log = LoggerFactory.getLogger(RuntimeConfigController.class); + + private final ResourceLoader resourceLoader; + private final ObjectMapper objectMapper; + private final @Nullable String webUiDir; + private final String backendUrl; + + /** Reads the bundled config through {@code resourceLoader}, or from {@code webUiDir} if set. */ + @Autowired + public RuntimeConfigController( + ResourceLoader resourceLoader, + ObjectMapper objectMapper, + @Value("${adk.web.ui.dir:#{null}}") @Nullable String webUiDir, + BackendUrl backendUrl) { + this.resourceLoader = resourceLoader; + this.objectMapper = objectMapper; + this.webUiDir = webUiDir; + this.backendUrl = backendUrl.value(); + } + + /** Serves the bundled config with {@code backendUrl} taken from configuration when set. */ + @GetMapping( + value = "/dev-ui/" + DevUiAssets.RUNTIME_CONFIG_PATH, + produces = MediaType.APPLICATION_JSON_VALUE) + public ResponseEntity> runtimeConfig() { + Map config = readBundledConfig(); + // Unset leaves a value the bundled document already carries, which is what used to be served. + if (backendUrl.isEmpty()) { + config.putIfAbsent("backendUrl", ""); + } else { + config.put("backendUrl", backendUrl); + } + // The bundled document can change on disk under adk.web.ui.dir, so do not let it be cached. + return ResponseEntity.ok().cacheControl(CacheControl.noStore()).body(config); + } + + /** + * Reads the bundled config, returning an empty map if the resource is missing or cannot be parsed + * so the endpoint still serves {@code backendUrl}. + */ + private Map readBundledConfig() { + Resource resource = + resourceLoader.getResource( + DevUiAssets.assetLocation(webUiDir, DevUiAssets.RUNTIME_CONFIG_PATH)); + if (!resource.exists()) { + log.debug("No bundled dev UI runtime config at {}; serving backendUrl only.", resource); + return new LinkedHashMap<>(); + } + try (InputStream in = resource.getInputStream()) { + Map parsed = objectMapper.readValue(in, new TypeReference<>() {}); + return parsed == null ? new LinkedHashMap<>() : new LinkedHashMap<>(parsed); + } catch (IOException e) { + log.warn("Could not read the bundled dev UI runtime config at {}.", resource, e); + return new LinkedHashMap<>(); + } + } +} diff --git a/dev/src/test/java/com/google/adk/web/AdkWebServerProxyRedirectTest.java b/dev/src/test/java/com/google/adk/web/AdkWebServerProxyRedirectTest.java index 2f4e3bfd5..3a323d431 100644 --- a/dev/src/test/java/com/google/adk/web/AdkWebServerProxyRedirectTest.java +++ b/dev/src/test/java/com/google/adk/web/AdkWebServerProxyRedirectTest.java @@ -29,14 +29,10 @@ import org.springframework.test.web.servlet.MockMvc; /** - * Spring's own forwarded-header support keeps working: with {@code - * server.forward-headers-strategy=framework} the operator opts in, {@code ForwardedHeaderFilter} - * turns the forwarded prefix into the request's context path, and the redirect picks it up. This - * server adds nothing here and reads no header itself; the test exists so enabling that Spring - * feature keeps behaving as it did. - * - *

The filter honours the standard {@code Forwarded} header as well as {@code X-Forwarded-*}, so - * both are covered. + * Verifies that when {@code adk.web.backend-url} is unset and {@code + * server.forward-headers-strategy=framework} is enabled, Spring's {@code ForwardedHeaderFilter} + * applies {@code Forwarded} and {@code X-Forwarded-*} headers to the entry redirects. When {@code + * adk.web.backend-url} is set, its path prefix takes precedence over {@code X-Forwarded-Prefix}. */ @SpringBootTest(properties = "server.forward-headers-strategy=framework") @AutoConfigureMockMvc diff --git a/dev/src/test/java/com/google/adk/web/BackendUrlRedirectTest.java b/dev/src/test/java/com/google/adk/web/BackendUrlRedirectTest.java new file mode 100644 index 000000000..ce6d61381 --- /dev/null +++ b/dev/src/test/java/com/google/adk/web/BackendUrlRedirectTest.java @@ -0,0 +1,169 @@ +/* + * 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.web; + +import static org.springframework.test.web.servlet.request.MockMvcRequestBuilders.get; +import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.redirectedUrl; +import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.status; + +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.boot.test.context.SpringBootTest; +import org.springframework.boot.webmvc.test.autoconfigure.AutoConfigureMockMvc; +import org.springframework.test.web.servlet.MockMvc; + +/** + * With {@code adk.web.backend-url} set, the entry redirect carries the gateway's path prefix on its + * own, so a deployment behind a path-stripping proxy needs nothing forwarded from the proxy. The + * property governs the path only: host and scheme still come from forwarded headers wherever the + * operator has enabled {@code server.forward-headers-strategy}. + */ +public class BackendUrlRedirectTest { + + @Nested + @SpringBootTest( + properties = { + "adk.web.backend-url=https://gw.example.com/my-app", + // Enable ForwardedHeaderFilter so X-Forwarded-* headers reach the servlet request. + "server.forward-headers-strategy=framework" + }) + @AutoConfigureMockMvc + class ConfiguredWithTheFilter { + + @ParameterizedTest + @ValueSource(strings = {"/", "/dev-ui"}) + public void devUiEntryPoints_shouldCarryTheConfiguredPrefix( + String path, @Autowired MockMvc mockMvc) throws Exception { + mockMvc + .perform(get(path)) + .andExpect(status().is3xxRedirection()) + .andExpect(redirectedUrl("/my-app/dev-ui/")); + } + + @ParameterizedTest + @ValueSource(strings = {"/", "/dev-ui"}) + public void devUiEntryPoints_shouldNotStackAForwardedPrefix( + String path, @Autowired MockMvc mockMvc) throws Exception { + // The configured prefix replaces the context path created by ForwardedHeaderFilter. + mockMvc + .perform(get(path).header("X-Forwarded-Prefix", "/evil")) + .andExpect(status().is3xxRedirection()) + .andExpect(redirectedUrl("http://localhost/my-app/dev-ui/")); + } + + @Test + public void devUiEntryPoint_forwardedHostAndProto_stillSetTheRedirectsOrigin( + @Autowired MockMvc mockMvc) throws Exception { + // The configured value supplies the path; host and scheme still come from the headers. + mockMvc + .perform( + get("/") + .header("X-Forwarded-Host", "gw.example.com") + .header("X-Forwarded-Proto", "https")) + .andExpect(status().is3xxRedirection()) + .andExpect(redirectedUrl("https://gw.example.com/my-app/dev-ui/")); + } + } + + @Nested + @SpringBootTest(properties = "adk.web.backend-url=https://gw.example.com/my-app") + @AutoConfigureMockMvc + class ConfiguredWithoutTheFilter { + + @Test + public void devUiEntryPoint_forwardedHeaders_reachNothingUnderTheShippedStrategy( + @Autowired MockMvc mockMvc) throws Exception { + // With forward-headers-strategy unset (default), forwarded headers are ignored. + mockMvc + .perform( + get("/") + .header("X-Forwarded-Prefix", "/evil") + .header("X-Forwarded-Host", "evil.example.com") + .header("Forwarded", "host=evil.example.com;proto=https")) + .andExpect(status().is3xxRedirection()) + .andExpect(redirectedUrl("/my-app/dev-ui/")); + } + } + + @Nested + @SpringBootTest(properties = "adk.web.backend-url=https://gw.example.com") + @AutoConfigureMockMvc + class ConfiguredWithNoPath { + + @ParameterizedTest + @ValueSource(strings = {"/", "/dev-ui"}) + public void devUiEntryPoints_shouldRedirectWithoutAPrefix( + String path, @Autowired MockMvc mockMvc) throws Exception { + // A host-only URL has an empty pathPrefix(), so the redirect stays context-relative. + mockMvc + .perform(get(path)) + .andExpect(status().is3xxRedirection()) + .andExpect(redirectedUrl("/dev-ui/")); + } + } + + @Nested + @SpringBootTest(properties = "adk.web.backend-url=/my-app") + @AutoConfigureMockMvc + class ConfiguredUnusably { + + @ParameterizedTest + @ValueSource(strings = {"/", "/dev-ui"}) + public void devUiEntryPoints_shouldRedirectWithoutAPrefix( + String path, @Autowired MockMvc mockMvc) throws Exception { + // A scheme-less value is unusable by the UI, so pathPrefix() is empty. + mockMvc + .perform(get(path)) + .andExpect(status().is3xxRedirection()) + .andExpect(redirectedUrl("/dev-ui/")); + } + } + + @Nested + @SpringBootTest + @AutoConfigureMockMvc + class UnconfiguredWithoutTheFilter { + + @ParameterizedTest + @ValueSource(strings = {"/", "/dev-ui"}) + public void devUiEntryPoints_shouldRedirectWithoutAPrefix( + String path, @Autowired MockMvc mockMvc) throws Exception { + mockMvc + .perform(get(path)) + .andExpect(status().is3xxRedirection()) + .andExpect(redirectedUrl("/dev-ui/")); + } + + @ParameterizedTest + @ValueSource(strings = {"/", "/dev-ui"}) + public void devUiEntryPoints_withForwardedHeaders_shouldStillRedirectWithoutAPrefix( + String path, @Autowired MockMvc mockMvc) throws Exception { + // By default forward-headers-strategy is unset, so forwarded headers are ignored. + mockMvc + .perform( + get(path) + .header("X-Forwarded-Prefix", "/evil") + .header("X-Forwarded-Host", "evil.example.com") + .header("Forwarded", "host=evil.example.com;proto=https")) + .andExpect(status().is3xxRedirection()) + .andExpect(redirectedUrl("/dev-ui/")); + } + } +} diff --git a/dev/src/test/java/com/google/adk/web/config/BackendUrlTest.java b/dev/src/test/java/com/google/adk/web/config/BackendUrlTest.java new file mode 100644 index 000000000..3b6d79762 --- /dev/null +++ b/dev/src/test/java/com/google/adk/web/config/BackendUrlTest.java @@ -0,0 +1,188 @@ +/* + * 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.web.config; + +import static com.google.common.truth.Truth.assertThat; + +import ch.qos.logback.classic.Level; +import ch.qos.logback.classic.Logger; +import ch.qos.logback.classic.spi.ILoggingEvent; +import ch.qos.logback.core.read.ListAppender; +import java.util.List; +import org.junit.jupiter.api.Test; +import org.slf4j.LoggerFactory; + +/** + * One reading of {@code adk.web.backend-url}. The value the dev UI is served and the prefix its + * entry redirect carries come from the same parse, so they cannot disagree. + */ +public class BackendUrlTest { + + private static final String CONFIGURED = "https://gw.example.com/my-app"; + + @Test + public void unset_isEmptyEverywhere() { + for (String in : new String[] {null, "", " "}) { + assertThat(BackendUrl.from(in).value()).isEmpty(); + assertThat(BackendUrl.from(in).pathPrefix()).isEmpty(); + } + } + + @Test + public void absoluteUrl_isServedAndSuppliesThePrefix() { + BackendUrl url = BackendUrl.from(CONFIGURED); + + assertThat(url.value()).isEqualTo(CONFIGURED); + assertThat(url.pathPrefix()).isEqualTo("/my-app"); + } + + @Test + public void trailingSlashes_areStripped() { + // The UI appends paths starting with a slash, so a trailing one would yield //run_live. + assertThat(BackendUrl.from(CONFIGURED + "/").value()).isEqualTo(CONFIGURED); + assertThat(BackendUrl.from(CONFIGURED + "///").pathPrefix()).isEqualTo("/my-app"); + assertThat(BackendUrl.from(" " + CONFIGURED + " ").value()).isEqualTo(CONFIGURED); + } + + @Test + public void hostWithoutPath_hasNoPrefix() { + assertThat(BackendUrl.from("https://gw.example.com").pathPrefix()).isEmpty(); + assertThat(BackendUrl.from("https://gw.example.com/").pathPrefix()).isEmpty(); + } + + @Test + public void percentEncoding_isKept() { + // This goes into a Location header, so decoding would re-encode wrongly and %2F would + // turn into a path separator. + assertThat(BackendUrl.from("https://gw.example.com/my%20app").pathPrefix()) + .isEqualTo("/my%20app"); + assertThat(BackendUrl.from("https://gw.example.com/a%2Fb").pathPrefix()).isEqualTo("/a%2Fb"); + } + + @Test + public void doubledSlash_doesNotBecomeAHost() { + // "//my-app/dev-ui/" is protocol-relative: a browser resolves it to the host "my-app". + assertThat(BackendUrl.from("https://gw.example.com//my-app").pathPrefix()).isEqualTo("/my-app"); + } + + @Test + public void unusableValue_suppliesNoPrefix() { + // Leave pathPrefix() empty when the UI cannot use the URL so redirect and UI behavior match. + for (String in : + new String[] { + "/my-app", + "HTTPS://gw.example.com/x", + "http://", + "gw.example.com", + "https://gw.example.com/a?q=1", + "https://user:pass@gw.example.com/my-app", + "https://gw.example.com/my app", + // No host at all, only a port: getRawAuthority() would call this usable. + "https://:8080/my-app", + // RFC-invalid, and unreachable anyway - Tomcat rejects such a Host header with a 400, + // and no CA will issue a certificate for one. + "https://gw_host.example.com/my-app" + }) { + assertThat(BackendUrl.from(in).pathPrefix()).isEmpty(); + } + } + + @Test + public void credentials_areNotLogged() { + String warning = warningsFor("https://user:pass@gw.example.com/my-app").get(0); + + assertThat(warning).doesNotContain("pass"); + assertThat(warning).contains("***@gw.example.com"); + } + + @Test + public void unusableValue_isStillServedButWarns() { + // Never silently discarded, because it is an explicit setting. + assertThat(BackendUrl.from("/my-app").value()).isEqualTo("/my-app"); + assertThat(BackendUrl.from("HTTPS://gw.example.com/x").value()) + .isEqualTo("HTTPS://gw.example.com/x"); + assertThat(BackendUrl.from("http://").value()).isEqualTo("http://"); + + assertThat(warningsFor("/my-app")).hasSize(1); + assertThat(warningsFor("/my-app").get(0)).contains("/my-app"); + assertThat(warningsFor("HTTPS://gw.example.com/x")).hasSize(1); + assertThat(warningsFor("http://")).hasSize(1); + } + + @Test + public void slashOnly_isNotEmptiedByTheSlashTrim() { + // Emptying this would make the served value fall back to whatever the bundled config says, + // silently losing an explicit setting. + assertThat(BackendUrl.from("/").value()).isEqualTo("/"); + assertThat(BackendUrl.from("/").pathPrefix()).isEmpty(); + } + + @Test + public void queryFragmentOrUserinfo_yieldsNoPrefixAndWarns() { + // The UI appends onto the value, so these would land in the middle of every request. + for (String in : + new String[] { + "https://gw.example.com/a/b?q=1", + "https://gw.example.com/a/b#f", + "https://user:pass@gw.example.com/my-app" + }) { + assertThat(BackendUrl.from(in).pathPrefix()).isEmpty(); + assertThat(warningsFor(in)).hasSize(1); + } + } + + @Test + public void valueWithoutALeadingSlashPath_yieldsNoPrefix() { + // "gw.example.com/dev-ui/" in a Location header would be resolved as a relative path. + assertThat(BackendUrl.from("gw.example.com").pathPrefix()).isEmpty(); + assertThat(warningsFor("gw.example.com")).hasSize(1); + } + + @Test + public void unparseableValue_yieldsNoPrefixAndWarns() { + String malformed = "https://gw.example.com/my app"; + + assertThat(BackendUrl.from(malformed).value()).isEqualTo(malformed); + assertThat(BackendUrl.from(malformed).pathPrefix()).isEmpty(); + assertThat(warningsFor(malformed)).hasSize(1); + assertThat(warningsFor(malformed).get(0)).contains("not a valid URI"); + } + + @Test + public void usableValue_doesNotWarn() { + assertThat(warningsFor(CONFIGURED)).isEmpty(); + assertThat(warningsFor(CONFIGURED + "/")).isEmpty(); + assertThat(warningsFor("")).isEmpty(); + } + + /** The WARN messages logged while interpreting {@code configured}. */ + private static List warningsFor(String configured) { + Logger logger = (Logger) LoggerFactory.getLogger(BackendUrl.class); + ListAppender appender = new ListAppender<>(); + appender.start(); + logger.addAppender(appender); + try { + BackendUrl unused = BackendUrl.from(configured); + } finally { + logger.detachAppender(appender); + } + return appender.list.stream() + .filter(event -> event.getLevel() == Level.WARN) + .map(ILoggingEvent::getFormattedMessage) + .toList(); + } +} diff --git a/dev/src/test/java/com/google/adk/web/config/DevUiAssetsTest.java b/dev/src/test/java/com/google/adk/web/config/DevUiAssetsTest.java index 668f75e57..46ccaf97c 100644 --- a/dev/src/test/java/com/google/adk/web/config/DevUiAssetsTest.java +++ b/dev/src/test/java/com/google/adk/web/config/DevUiAssetsTest.java @@ -42,4 +42,10 @@ public void assetRoot_configuredDir_becomesAFileUrl(String webUiDir) { public void assetRoot_windowsSeparators_areNormalized() { assertThat(DevUiAssets.assetRoot("C:\\srv\\ui")).isEqualTo("file:C:/srv/ui/"); } + + @Test + public void assetLocation_appendsToTheRoot() { + assertThat(DevUiAssets.assetLocation(null, DevUiAssets.RUNTIME_CONFIG_PATH)) + .isEqualTo("classpath:/browser/assets/config/runtime-config.json"); + } } diff --git a/dev/src/test/java/com/google/adk/web/controller/RuntimeConfigControllerMergeTest.java b/dev/src/test/java/com/google/adk/web/controller/RuntimeConfigControllerMergeTest.java new file mode 100644 index 000000000..a01dd9462 --- /dev/null +++ b/dev/src/test/java/com/google/adk/web/controller/RuntimeConfigControllerMergeTest.java @@ -0,0 +1,123 @@ +/* + * 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.web.controller; + +import static com.google.common.truth.Truth.assertThat; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.adk.web.config.BackendUrl; +import java.nio.charset.StandardCharsets; +import java.util.Map; +import org.junit.jupiter.api.Test; +import org.springframework.core.io.ByteArrayResource; +import org.springframework.core.io.DescriptiveResource; +import org.springframework.core.io.Resource; +import org.springframework.core.io.ResourceLoader; + +/** + * The served runtime config merges into the bundled document rather than replacing it, so keys the + * dev UI bundle gains later are not silently dropped. {@code adk.web.backend-url} overrides one + * key; an unset value leaves the bundled document exactly as it was served before. + */ +public class RuntimeConfigControllerMergeTest { + + private static final String CONFIGURED = "https://gw.example.com/my-app"; + + @Test + public void runtimeConfig_shouldPreserveOtherBundledKeys() { + String bundled = "{\"backendUrl\":\"\",\"telemetry\":null,\"logo\":{\"text\":\"x\"}}"; + RuntimeConfigController controller = controllerFor(bundled, CONFIGURED); + + Map config = controller.runtimeConfig().getBody(); + + assertThat(config).containsEntry("backendUrl", CONFIGURED); + assertThat(config).containsKey("telemetry"); + assertThat(config).containsEntry("logo", Map.of("text", "x")); + } + + @Test + public void runtimeConfig_shouldOverrideTheBundledBackendUrl() { + RuntimeConfigController controller = + controllerFor("{\"backendUrl\":\"http://stale\"}", CONFIGURED); + + assertThat(controller.runtimeConfig().getBody()).containsEntry("backendUrl", CONFIGURED); + } + + @Test + public void runtimeConfig_propertyUnset_shouldKeepTheBundledBackendUrl() { + // The static handler served this file verbatim, so a hand-set value survived; it still does. + RuntimeConfigController controller = + controllerFor("{\"backendUrl\":\"http://elsewhere:9000\"}"); + + assertThat(controller.runtimeConfig().getBody()) + .containsExactly("backendUrl", "http://elsewhere:9000"); + } + + @Test + public void runtimeConfig_propertyUnsetAndNoBundledKey_shouldStillReportBackendUrl() { + RuntimeConfigController controller = controllerFor("{\"telemetry\":null}"); + + assertThat(controller.runtimeConfig().getBody()).containsEntry("backendUrl", ""); + } + + @Test + public void runtimeConfig_bundledFileMissing_shouldStillServe() { + RuntimeConfigController controller = controllerFor(null, CONFIGURED); + + assertThat(controller.runtimeConfig().getBody()).containsExactly("backendUrl", CONFIGURED); + } + + @Test + public void runtimeConfig_bundledFileMalformed_shouldStillServe() { + RuntimeConfigController controller = controllerFor("not json at all", CONFIGURED); + + assertThat(controller.runtimeConfig().getBody()).containsExactly("backendUrl", CONFIGURED); + } + + @Test + public void runtimeConfig_bundledFileNotAnObject_shouldStillServe() { + // A config that is not a JSON object is ignored rather than failing the request. + RuntimeConfigController controller = controllerFor("[1, 2, 3]", CONFIGURED); + + assertThat(controller.runtimeConfig().getBody()).containsExactly("backendUrl", CONFIGURED); + } + + /** A controller with {@code adk.web.backend-url} unset. */ + private static RuntimeConfigController controllerFor(String body) { + return controllerFor(body, ""); + } + + /** A controller whose bundled config is {@code body}, or absent when {@code body} is null. */ + private static RuntimeConfigController controllerFor(String body, String backendUrl) { + ResourceLoader loader = + new ResourceLoader() { + @Override + public Resource getResource(String location) { + return body == null + ? new DescriptiveResource("absent") + : new ByteArrayResource(body.getBytes(StandardCharsets.UTF_8)); + } + + @Override + public ClassLoader getClassLoader() { + return RuntimeConfigControllerMergeTest.class.getClassLoader(); + } + }; + return new RuntimeConfigController( + loader, new ObjectMapper(), null, BackendUrl.from(backendUrl)); + } +} diff --git a/dev/src/test/java/com/google/adk/web/controller/RuntimeConfigControllerTest.java b/dev/src/test/java/com/google/adk/web/controller/RuntimeConfigControllerTest.java new file mode 100644 index 000000000..4356fc2a4 --- /dev/null +++ b/dev/src/test/java/com/google/adk/web/controller/RuntimeConfigControllerTest.java @@ -0,0 +1,121 @@ +/* + * 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.web.controller; + +import static org.springframework.test.web.servlet.request.MockMvcRequestBuilders.get; +import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.content; +import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.header; +import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.jsonPath; +import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.status; + +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.boot.test.context.SpringBootTest; +import org.springframework.boot.webmvc.test.autoconfigure.AutoConfigureMockMvc; +import org.springframework.test.web.servlet.MockMvc; + +/** + * The dev UI learns where its backend lives from {@code adk.web.backend-url}, so its API calls + * reach a server published under a path prefix. Nothing is read from the request, so a client + * cannot influence the value it is served. + */ +public class RuntimeConfigControllerTest { + + private static final String CONFIG = "/dev-ui/assets/config/runtime-config.json"; + + @Nested + @SpringBootTest( + properties = { + "adk.web.backend-url=https://gw.example.com/my-app", + // Enable ForwardedHeaderFilter so X-Forwarded-* headers reach the servlet request. + "server.forward-headers-strategy=framework" + }) + @AutoConfigureMockMvc + class ConfiguredWithTheFilter { + + @Test + public void runtimeConfig_shouldReportTheConfiguredBackendUrl(@Autowired MockMvc mockMvc) + throws Exception { + mockMvc + .perform(get(CONFIG)) + .andExpect(status().isOk()) + .andExpect(jsonPath("$.backendUrl").value("https://gw.example.com/my-app")) + .andExpect(header().string("Cache-Control", "no-store")); + } + + @Test + public void runtimeConfig_forwardedHeaders_shouldNotChangeTheValue(@Autowired MockMvc mockMvc) + throws Exception { + mockMvc + .perform( + get(CONFIG) + .header("X-Forwarded-Prefix", "/evil") + .header("X-Forwarded-Host", "evil.example.com") + .header("X-Forwarded-Proto", "https")) + .andExpect(status().isOk()) + .andExpect(jsonPath("$.backendUrl").value("https://gw.example.com/my-app")); + } + } + + @Nested + @SpringBootTest + @AutoConfigureMockMvc + class UnconfiguredWithoutTheFilter { + + @Test + public void runtimeConfig_forwardedHeaders_areNotReadFromTheRequest(@Autowired MockMvc mockMvc) + throws Exception { + mockMvc + .perform( + get(CONFIG) + .header("X-Forwarded-Prefix", "/evil") + .header("X-Forwarded-Host", "evil.example.com") + .header("Forwarded", "host=evil.example.com;proto=https")) + .andExpect(status().isOk()) + .andExpect(jsonPath("$.backendUrl").value("")); + } + } + + @Nested + @SpringBootTest(properties = "server.forward-headers-strategy=framework") + @AutoConfigureMockMvc + class UnconfiguredWithTheFilter { + + @Test + public void runtimeConfig_shouldServeWhatTheBundledFileSaid(@Autowired MockMvc mockMvc) + throws Exception { + mockMvc + .perform(get(CONFIG)) + .andExpect(status().isOk()) + // Compact JSON proves the controller answered rather than the pretty-printed static file. + .andExpect(content().string("{\"backendUrl\":\"\"}")); + } + + @Test + public void runtimeConfig_forwardedHeaders_shouldStillReportEmpty(@Autowired MockMvc mockMvc) + throws Exception { + mockMvc + .perform( + get(CONFIG) + .header("X-Forwarded-Prefix", "/evil") + .header("X-Forwarded-Host", "evil.example.com")) + .andExpect(status().isOk()) + .andExpect(jsonPath("$.backendUrl").value("")); + } + } +} diff --git a/tokt/src/main/kotlin/com/google/adk/tokt/KtRunnerToJava.kt b/tokt/src/main/kotlin/com/google/adk/tokt/KtRunnerToJava.kt index c141d184e..0467ffec1 100644 --- a/tokt/src/main/kotlin/com/google/adk/tokt/KtRunnerToJava.kt +++ b/tokt/src/main/kotlin/com/google/adk/tokt/KtRunnerToJava.kt @@ -18,10 +18,13 @@ package com.google.adk.tokt import com.google.adk.agents.LiveRequestQueue import com.google.adk.agents.RunConfig as JavaRunConfig +import com.google.adk.annotations.Experimental import com.google.adk.artifacts.BaseArtifactService as JavaArtifactService import com.google.adk.artifacts.InMemoryArtifactService as JavaInMemoryArtifactService import com.google.adk.events.Event as JavaEvent import com.google.adk.kt.runners.Runner as KtRunner +import com.google.adk.kt.types.Content as KtContent +import com.google.adk.kt.types.Role as KtRole import com.google.adk.plugins.PluginManager as JavaPluginManager import com.google.adk.runner.Runner as JavaRunner import com.google.adk.sessions.Session as JavaSession @@ -78,6 +81,29 @@ internal class KtRunnerToJava(private val ktRunner: KtRunner, dispatcher: Corout newMessage: GenaiContent, runConfig: JavaRunConfig, stateDelta: MutableMap?, + ): Flowable = + runAsync( + userId, + sessionId, + invocationId = null, + newMessage = newMessage, + runConfig = runConfig, + stateDelta = stateDelta, + ) + + /** + * Resumes on the Kotlin engine with its semantics instead of the Java runner's. A non-resumable + * app runs [newMessage] as a new invocation under [invocationId], as Python 2.x does, instead of + * throwing. ADK Kotlin 1.1.0 and earlier drop [stateDelta] when [newMessage] is null. + */ + @Experimental + override fun runAsync( + userId: String, + sessionId: String, + invocationId: String?, + newMessage: GenaiContent?, + runConfig: JavaRunConfig, + stateDelta: MutableMap?, ): Flowable { // Defer so a request-conversion or RunConfig rejection surfaces via onError, not at the call // site - matching the base Runner and this class's own runLive. @@ -87,7 +113,8 @@ internal class KtRunnerToJava(private val ktRunner: KtRunner, dispatcher: Corout .runAsync( userId = userId, sessionId = sessionId, - newMessage = ContentCodec.fromJava(newMessage), + invocationId = invocationId, + newMessage = newMessage?.let { withUserRole(ContentCodec.fromJava(it)) }, // Translate the Java REMOVED sentinel so a caller-passed deletion deletes the key // rather than storing it as a value (identity-matched by the engine). stateDelta = stateDelta?.let { stateDeltaFromJava(it) }, @@ -146,4 +173,11 @@ internal class KtRunnerToJava(private val ktRunner: KtRunner, dispatcher: Corout private fun liveUnsupported() = UnsupportedOperationException("Live mode is not supported when running on the Kotlin engine.") + + /** + * Defaults a role-less message to the user role, as Python does, since model requests omit + * role-less content. + */ + private fun withUserRole(content: KtContent): KtContent = + if (content.role.isNullOrEmpty()) content.copy(role = KtRole.USER) else content } diff --git a/tokt/src/test/kotlin/com/google/adk/tokt/KtRunnerInteropTest.kt b/tokt/src/test/kotlin/com/google/adk/tokt/KtRunnerInteropTest.kt index b1861703e..f83fb022c 100644 --- a/tokt/src/test/kotlin/com/google/adk/tokt/KtRunnerInteropTest.kt +++ b/tokt/src/test/kotlin/com/google/adk/tokt/KtRunnerInteropTest.kt @@ -29,6 +29,7 @@ import com.google.adk.events.EventCompaction as JavaEventCompaction import com.google.adk.kt.agents.BaseAgent as KtBaseAgent import com.google.adk.kt.agents.InvocationContext as KtInvocationContext import com.google.adk.kt.agents.LlmAgent as KtLlmAgent +import com.google.adk.kt.agents.ResumabilityConfig as KtResumabilityConfig import com.google.adk.kt.agents.RunConfig as KtRunConfig import com.google.adk.kt.agents.StreamingMode as KtStreamingMode import com.google.adk.kt.agents.TypedData @@ -89,6 +90,7 @@ import com.google.adk.sessions.GetSessionConfig as JavaGetSessionConfig import com.google.adk.sessions.InMemorySessionService as JavaInMemorySessionService import com.google.adk.sessions.ListEventsResponse as JavaListEventsResponse import com.google.adk.sessions.Session as JavaSession +import com.google.adk.sessions.SessionKey as JavaSessionKey import com.google.adk.sessions.State as JavaState import com.google.adk.tokt.adapters.reconcileActionsToKt import com.google.adk.tokt.codecs.EventCodec @@ -113,6 +115,7 @@ import com.google.genai.types.ExecutableCode as GenaiExecutableCode import com.google.genai.types.FinishReason as GenaiFinishReason import com.google.genai.types.FunctionCall as GenaiFunctionCall import com.google.genai.types.FunctionDeclaration as GenaiFunctionDeclaration +import com.google.genai.types.FunctionResponse as GenaiFunctionResponse import com.google.genai.types.GenerateContentResponseUsageMetadata as GenaiUsageMetadata import com.google.genai.types.GroundingChunk as GenaiGroundingChunk import com.google.genai.types.GroundingChunkMaps as GenaiGroundingChunkMaps @@ -145,6 +148,7 @@ import java.util.concurrent.atomic.AtomicInteger import java.util.concurrent.atomic.AtomicReference import kotlin.jvm.optionals.getOrNull import kotlin.test.Test +import kotlin.test.assertContains import kotlin.test.assertContentEquals import kotlin.test.assertEquals import kotlin.test.assertFailsWith @@ -278,6 +282,19 @@ class KtRunnerInteropTest { } } + /** A long-running Java tool that returns only a pending status. */ + private class JavaPendingTool : + JavaBaseTool("java_pending", "long-running", /* isLongRunning= */ true) { + override fun declaration(): Optional = + Optional.of(GenaiFunctionDeclaration.builder().name("java_pending").build()) + + @JvmSuppressWildcards + override fun runAsync( + args: Map, + toolContext: JavaToolContext, + ): Single> = Single.just(mapOf("status" to "pending")) + } + /** * A Java tool that mutates its [JavaToolContext.actions] in place (not via `setActions`) to set * every control-flow signal, exercising the live [KtEventActionsToJavaView] write-through. @@ -701,6 +718,29 @@ class KtRunnerInteropTest { } } + /** A Kotlin runner that records each runAsync call once its stream is collected. */ + private class RecordingKtRunner(private val delegate: KtRunner) : KtRunner by delegate { + data class RunAsyncCall( + val userId: String, + val sessionId: String, + val invocationId: String?, + val newMessage: KtContent?, + ) + + val runAsyncCalls = mutableListOf() + + override fun runAsync( + userId: String, + sessionId: String, + invocationId: String?, + newMessage: KtContent?, + stateDelta: Map?, + runConfig: KtRunConfig?, + ): Flow = flow { + runAsyncCalls.add(RunAsyncCall(userId, sessionId, invocationId, newMessage)) + } + } + @Test fun javaAdkToKt_convertsEntireCollections() { // Tools: order preserved, each Java tool wrapped as a Kotlin tool. @@ -1159,6 +1199,307 @@ class KtRunnerInteropTest { ) } + @Test + fun asJavaRunner_resumeOverload_resumesThePausedInvocation() { + // A resumable app pauses on a confirmation; the resume overload continues that invocation. + val ktRunner = + KtInMemoryRunner( + app = + KtApp( + appName = "app", + rootAgent = + KtLlmAgent( + name = "a", + model = + JavaAdkToKt.asKtModel( + SequentialJavaModel( + listOf( + modelFunctionCall("java_confirm", emptyMap()), + modelText("confirmed done"), + ) + ) + ), + tools = listOf(JavaAdkToKt.asKtTool(JavaConfirmTool())), + ), + resumabilityConfig = KtResumabilityConfig(isResumable = true), + ) + ) + val javaRunner = KotlinAdkToJava.asJavaRunner(ktRunner) + val autoCreate = JavaRunConfig.builder().autoCreateSession(true).build() + + val turn1 = + javaRunner + .runAsync( + "u", + "s", + GenaiContent.builder().role("user").parts(GenaiPart.fromText("go")).build(), + autoCreate, + ) + .toList() + .blockingGet() + val (pausedEvent, confirmationCall) = + turn1 + .flatMap { event -> event.functionCalls().map { event to it } } + .first { (_, call) -> call.name().getOrNull() == "adk_request_confirmation" } + val pausedInvocationId = pausedEvent.invocationId() + val approval = + GenaiContent.builder() + .role("user") + .parts( + GenaiPart.builder() + .functionResponse( + GenaiFunctionResponse.builder() + .name("adk_request_confirmation") + .id(confirmationCall.id().get()) + .response(mapOf("confirmed" to true)) + .build() + ) + .build() + ) + .build() + + // The engine infers the paused invocation from the approval's function-response id. + val turn2 = + javaRunner + .runAsync("u", "s", /* invocationId= */ null, approval, autoCreate, /* stateDelta= */ null) + .toList() + .blockingGet() + + val (resultEvent, toolResponse) = + turn2 + .flatMap { event -> event.functionResponses().map { event to it } } + .first { (_, response) -> response.name().getOrNull() == "java_confirm" } + assertEquals( + "confirmed", + toolResponse.response().getOrNull()?.get("status"), + "the resumed run should deliver the approval to the tool", + ) + assertEquals( + pausedInvocationId, + resultEvent.invocationId(), + "turn 2 should continue the paused invocation, not start a new one", + ) + + // Resuming the finished invocation by id alone is a no-op rather than an error. + val turn3 = + javaRunner + .runAsync( + "u", + "s", + pausedInvocationId, + /* newMessage= */ null, + autoCreate, + /* stateDelta= */ null, + ) + .toList() + .blockingGet() + assertTrue(turn3.isEmpty(), "resuming a finished invocation should emit nothing") + } + + @Test + fun asJavaRunner_resumeOverload_roleLessResponse_reachesTheModelAsAUserTurn() { + val model = + SequentialJavaModel(listOf(modelFunctionCall("java_pending", emptyMap()), modelText("done"))) + val ktRunner = + KtInMemoryRunner( + app = + KtApp( + appName = "app", + rootAgent = + KtLlmAgent( + name = "a", + model = JavaAdkToKt.asKtModel(model), + tools = listOf(JavaAdkToKt.asKtTool(JavaPendingTool())), + ), + resumabilityConfig = KtResumabilityConfig(isResumable = true), + ) + ) + val javaRunner = KotlinAdkToJava.asJavaRunner(ktRunner) + val autoCreate = JavaRunConfig.builder().autoCreateSession(true).build() + val turn1 = + javaRunner + .runAsync( + "u", + "s", + GenaiContent.builder().role("user").parts(GenaiPart.fromText("go")).build(), + autoCreate, + ) + .toList() + .blockingGet() + val call = + turn1.flatMap { it.functionCalls() }.first { it.name().getOrNull() == "java_pending" } + val roleLessResponse = + GenaiContent.builder() + .parts( + GenaiPart.builder() + .functionResponse( + GenaiFunctionResponse.builder() + .name("java_pending") + .id(call.id().get()) + .response(mapOf("status" to "done")) + .build() + ) + .build() + ) + .build() + + val unusedResumeEvents = + javaRunner + .runAsync( + "u", + "s", + turn1.first().invocationId(), + roleLessResponse, + autoCreate, + /* stateDelta= */ null, + ) + .toList() + .blockingGet() + + val sentResponses = + model.requests + .last() + .contents() + .filter { it.role().getOrNull() == "user" } + .flatMap { it.parts().getOrNull().orEmpty() } + .mapNotNull { it.functionResponse().getOrNull()?.response()?.getOrNull() } + assertTrue( + mapOf("status" to "done") in sentResponses, + "the role-less response should reach the model as a user turn", + ) + } + + @Test + fun asJavaRunner_resumeOverload_nonResumableApp_startsNewInvocationWithCallerId() { + // The Java Runner throws for a non-resumable app; the Kotlin engine, like Python, runs it. + val javaRunner = + KotlinAdkToJava.asJavaRunner( + KtInMemoryRunner( + agent = + KtLlmAgent( + name = "a", + model = JavaAdkToKt.asKtModel(SequentialJavaModel(listOf(modelText("done")))), + ), + appName = "app", + ) + ) + + val events = + javaRunner + .runAsync( + "u", + "s", + "caller-id", + GenaiContent.builder().role("user").parts(GenaiPart.fromText("hi")).build(), + JavaRunConfig.builder().autoCreateSession(true).build(), + /* stateDelta= */ null, + ) + .toList() + .blockingGet() + + assertTrue(events.isNotEmpty(), "the run should start a new invocation instead of throwing") + assertTrue( + events.all { it.invocationId() == "caller-id" }, + "the new invocation should use the caller's invocation id", + ) + } + + @Test + fun asJavaRunner_resumeOverload_nonResumableApp_nullMessage_throws() { + val javaRunner = + KotlinAdkToJava.asJavaRunner( + KtInMemoryRunner( + agent = + KtLlmAgent(name = "a", model = JavaAdkToKt.asKtModel(SequentialJavaModel(emptyList()))), + appName = "app", + ) + ) + + val error = + assertFailsWith { + javaRunner + .runAsync( + "u", + "s", + "caller-id", + /* newMessage= */ null, + JavaRunConfig.builder().autoCreateSession(true).build(), + /* stateDelta= */ null, + ) + .toList() + .blockingGet() + } + assertContains(error.message.orEmpty(), "not resumable") + } + + @Test + fun asJavaRunner_everyRunAsyncOverload_routesToKotlinEngine() { + // Every public runAsync overload must reach the Kotlin engine, not the base Runner. + val recording = + RecordingKtRunner( + KtInMemoryRunner( + agent = + KtLlmAgent(name = "a", model = JavaAdkToKt.asKtModel(SequentialJavaModel(emptyList()))), + appName = "app", + ) + ) + val javaRunner = KotlinAdkToJava.asJavaRunner(recording) + val key = JavaSessionKey("app", "u", "s") + // Pre-create the session for overloads whose default RunConfig disables autoCreateSession. + javaRunner.sessionService().createSession(key).ignoreElement().blockingAwait() + + val msg = GenaiContent.builder().role("user").parts(GenaiPart.fromText("hi")).build() + val autoCreate = JavaRunConfig.builder().autoCreateSession(true).build() + + javaRunner.runAsync("u", "s", msg, autoCreate).ignoreElements().blockingAwait() + javaRunner + .runAsync("u", "s", msg, autoCreate, /* stateDelta= */ null) + .ignoreElements() + .blockingAwait() + javaRunner.runAsync("u", "s", msg).ignoreElements().blockingAwait() + javaRunner + .runAsync(key, msg, autoCreate, /* stateDelta= */ null) + .ignoreElements() + .blockingAwait() + javaRunner.runAsync(key, msg, autoCreate).ignoreElements().blockingAwait() + javaRunner.runAsync(key, msg).ignoreElements().blockingAwait() + javaRunner + .runAsync("u", "s", "invocation-xyz", msg, autoCreate, /* stateDelta= */ null) + .ignoreElements() + .blockingAwait() + javaRunner + .runAsync( + "u", + "s", + "invocation-xyz", + /* newMessage= */ null, + autoCreate, + /* stateDelta= */ null, + ) + .ignoreElements() + .blockingAwait() + + assertEquals( + 8, + recording.runAsyncCalls.size, + "each runAsync call should reach the Kotlin engine exactly once", + ) + assertTrue( + recording.runAsyncCalls.all { it.userId == "u" && it.sessionId == "s" }, + "every overload should pass userId and sessionId through", + ) + assertEquals( + listOf(null, null, null, null, null, null, "invocation-xyz", "invocation-xyz"), + recording.runAsyncCalls.map { it.invocationId }, + "only the resume overload should carry an invocationId, unchanged", + ) + assertEquals( + listOf(true, true, true, true, true, true, true, false), + recording.runAsyncCalls.map { it.newMessage != null }, + "a null message on the resume overload should reach the engine as null", + ) + } + @Test fun asJavaRunner_nonDefaultRunConfig_reachesTheKotlinEngine() { // A Java RunConfig set through the Java API must map onto the Kotlin engine; the plugin reads @@ -1209,25 +1550,28 @@ class KtRunnerInteropTest { appName = "app", ) val javaRunner = KotlinAdkToJava.asJavaRunner(ktRunner) + // Checks the adapter's own error, not the one the inspection-only agent view would raise. + fun assertLiveUnsupported(run: () -> Flowable) { + val error = assertFailsWith { run().toList().blockingGet() } + assertEquals("Live mode is not supported when running on the Kotlin engine.", error.message) + } // Live mode is not bridged; it must fail loudly through the stream (like the base Runner), not // throw eagerly at the call site. - assertFailsWith { - javaRunner - .runLive("u", "s", LiveRequestQueue(), JavaRunConfig.builder().build()) - .toList() - .blockingGet() + assertLiveUnsupported { + javaRunner.runLive("u", "s", LiveRequestQueue(), JavaRunConfig.builder().build()) } // agent() returns an inspection-only view of the Kotlin agent - readable, but not runnable. assertEquals("a", javaRunner.agent().name(), "agent() should expose the Kotlin agent's name") - // The session-based runLive overload must fail the same way, through the stream. - val session = javaRunner.sessionService().createSession("app", "u", null, "s").blockingGet() - assertFailsWith { - javaRunner - .runLive(session, LiveRequestQueue(), JavaRunConfig.builder().build()) - .toList() - .blockingGet() + // The session-based and SessionKey runLive overloads must fail the same way. + val key = JavaSessionKey("app", "u", "s") + val session = javaRunner.sessionService().createSession(key).blockingGet() + assertLiveUnsupported { + javaRunner.runLive(session, LiveRequestQueue(), JavaRunConfig.builder().build()) + } + assertLiveUnsupported { + javaRunner.runLive(key, LiveRequestQueue(), JavaRunConfig.builder().build()) } }