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 4e17780f8..b14e5abe1 100644 --- a/core/src/main/java/com/google/adk/agents/BaseAgent.java +++ b/core/src/main/java/com/google/adk/agents/BaseAgent.java @@ -276,7 +276,8 @@ public ImmutableList canonicalAfterAgentCallbacks( private InvocationContext createInvocationContext( InvocationContext parentContext, boolean keepBranch) { InvocationContext.Builder builder = parentContext.toBuilder(); - builder.agent(this); + // An agent's context runs that agent, not the node that its parent context ran. + builder.agent(this).node(null); if (!keepBranch) { // Check for branch to be truthy (not None, not empty string), parentContext 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 1a6f1786c..3a761dedc 100644 --- a/core/src/main/java/com/google/adk/agents/InvocationContext.java +++ b/core/src/main/java/com/google/adk/agents/InvocationContext.java @@ -33,6 +33,7 @@ import com.google.adk.sessions.BaseSessionService; import com.google.adk.sessions.Session; import com.google.adk.summarizer.EventsCompactionConfig; +import com.google.adk.workflow.Node; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; @@ -84,6 +85,7 @@ public class InvocationContext { @Nullable private String branch; private BaseAgent agent; + private final @Nullable Node node; private boolean endInvocation; protected InvocationContext(Builder builder) { @@ -97,6 +99,7 @@ protected InvocationContext(Builder builder) { // Resolved here so subclasses that construct from a builder without build() also get an id. this.invocationId = builder.resolveInvocationId(); this.agent = builder.agent; + this.node = builder.node; this.session = builder.session; this.userContent = builder.userContent; this.runConfig = builder.runConfig; @@ -182,6 +185,15 @@ public BaseAgent agent() { return agent; } + /** + * Returns the workflow node this invocation runs, or empty when it runs an agent. When it runs a + * node, {@link #agent()} returns a {@link NodeViewAgent} that stands in for it. + */ + @Experimental + public Optional node() { + return Optional.ofNullable(node); + } + /** Returns the session associated with this invocation. */ public Session session() { return session; @@ -653,6 +665,7 @@ private Builder(InvocationContext context) { this.invocationId = context.invocationId; this.invocationIdSet = true; this.agent = context.agent; + this.node = context.node; this.session = context.session; this.userContent = context.userContent; this.runConfig = context.runConfig; @@ -684,6 +697,7 @@ private Builder(InvocationContext context) { // Distinguishes "never set" (an id is minted from the UuidProvider) from an explicit null. private boolean invocationIdSet; private BaseAgent agent; + private @Nullable Node node; private Session session; @Nullable private Content userContent = null; private RunConfig runConfig = RunConfig.builder().build(); @@ -796,6 +810,17 @@ public Builder agent(BaseAgent agent) { return this; } + /** + * Sets the workflow node the invocation runs, or null when it runs an agent. When the + * invocation runs a node, its agent is a {@link NodeViewAgent} that stands in for the node. + */ + @CanIgnoreReturnValue + @Experimental + public Builder node(@Nullable Node node) { + this.node = node; + return this; + } + /** * Sets the session associated with this invocation. * @@ -968,6 +993,9 @@ private static void validate(Builder builder) { if (builder.agent == null) { throw new IllegalStateException("Agent must be set."); } + if (builder.agent instanceof NodeViewAgent && builder.node == null) { + throw new IllegalStateException("A node view agent requires the node it stands in for."); + } if (builder.session == null) { throw new IllegalStateException("Session must be set."); } @@ -994,6 +1022,7 @@ public boolean equals(Object o) { && Objects.equals(branch, that.branch) && Objects.equals(invocationId, that.invocationId) && Objects.equals(agent, that.agent) + && Objects.equals(node, that.node) && Objects.equals(session, that.session) && Objects.equals(userContent, that.userContent) && Objects.equals(runConfig, that.runConfig) @@ -1019,6 +1048,7 @@ public int hashCode() { branch, invocationId, agent, + node, session, userContent, runConfig, diff --git a/core/src/main/java/com/google/adk/agents/NodeViewAgent.java b/core/src/main/java/com/google/adk/agents/NodeViewAgent.java new file mode 100644 index 000000000..3fd733f53 --- /dev/null +++ b/core/src/main/java/com/google/adk/agents/NodeViewAgent.java @@ -0,0 +1,83 @@ +/* + * 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.agents; + +import static com.google.common.base.Preconditions.checkArgument; + +import com.google.adk.annotations.Experimental; +import com.google.adk.events.Event; +import com.google.adk.workflow.Node; +import io.reactivex.rxjava3.core.Flowable; + +/** + * A read-only agent that stands in for a non-agent workflow node so code reading the agent of an + * app, runner, or invocation keeps working when they run a node. It never runs; the runner runs the + * node itself. + */ +@Experimental +public final class NodeViewAgent extends BaseAgent { + + private final Node node; + + /** + * Creates an agent that stands in for {@code node}, using its name and description. + * + * @throws IllegalArgumentException if {@code node} is an agent, or if its name is not a valid + * agent name. + */ + public NodeViewAgent(Node node) { + super( + node.name(), + node.description(), + /* subAgents= */ null, + /* beforeAgentCallback= */ null, + /* afterAgentCallback= */ null); + checkArgument(!(node instanceof BaseAgent), "An agent needs no node view."); + this.node = node; + } + + /** Returns the node this agent stands in for. */ + public Node node() { + return node; + } + + // Callers get this error instead of the context check in BaseAgent.runAsync and runLive. + @Override + public Flowable runAsync(InvocationContext parentContext) { + return Flowable.error(neverRuns()); + } + + @Override + public Flowable runLive(InvocationContext parentContext) { + return Flowable.error(neverRuns()); + } + + @Override + protected Flowable runAsyncImpl(InvocationContext invocationContext) { + return Flowable.error(neverRuns()); + } + + @Override + protected Flowable runLiveImpl(InvocationContext invocationContext) { + return Flowable.error(neverRuns()); + } + + private UnsupportedOperationException neverRuns() { + return new UnsupportedOperationException( + "A node view agent never runs; it stands in for node '" + name() + "'."); + } +} 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 4d2d1b3d9..db1dc9b7a 100644 --- a/core/src/main/java/com/google/adk/apps/App.java +++ b/core/src/main/java/com/google/adk/apps/App.java @@ -16,15 +16,20 @@ package com.google.adk.apps; +import static com.google.common.base.Preconditions.checkNotNull; + import com.google.adk.agents.BaseAgent; import com.google.adk.agents.ContextCacheConfig; +import com.google.adk.agents.NodeViewAgent; 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.adk.workflow.Node; import com.google.common.collect.ImmutableList; import com.google.errorprone.annotations.CanIgnoreReturnValue; import java.util.List; +import java.util.Optional; import java.util.regex.Pattern; import org.jspecify.annotations.Nullable; @@ -41,6 +46,7 @@ public class App { private final String name; private final BaseAgent rootAgent; + private final @Nullable Node rootNode; private final ImmutableList plugins; private final @Nullable EventsCompactionConfig eventsCompactionConfig; private final @Nullable ContextCacheConfig contextCacheConfig; @@ -49,12 +55,14 @@ public class App { private App( String name, BaseAgent rootAgent, + @Nullable Node rootNode, List plugins, @Nullable EventsCompactionConfig eventsCompactionConfig, @Nullable ContextCacheConfig contextCacheConfig, @Nullable ResumabilityConfig resumabilityConfig) { this.name = name; this.rootAgent = rootAgent; + this.rootNode = rootNode; this.plugins = ImmutableList.copyOf(plugins); this.eventsCompactionConfig = eventsCompactionConfig; this.contextCacheConfig = contextCacheConfig; @@ -69,6 +77,15 @@ public BaseAgent rootAgent() { return rootAgent; } + /** + * Returns the workflow node the app runs, or empty when it runs an agent tree. When it runs a + * node, {@link #rootAgent()} returns a {@link NodeViewAgent} that stands in for it. + */ + @Experimental + public Optional rootNode() { + return Optional.ofNullable(rootNode); + } + public List plugins() { return plugins; } @@ -92,6 +109,7 @@ public ContextCacheConfig contextCacheConfig() { public static class Builder { private String name; private BaseAgent rootAgent; + private @Nullable Node rootNode; private List plugins = ImmutableList.of(); @Nullable private EventsCompactionConfig eventsCompactionConfig; @Nullable private ContextCacheConfig contextCacheConfig; @@ -109,6 +127,17 @@ public Builder rootAgent(BaseAgent rootAgent) { return this; } + /** + * Sets the workflow node the app runs in place of a root agent. If the given node is an agent, + * it becomes the root agent. + */ + @CanIgnoreReturnValue + @Experimental + public Builder rootNode(Node rootNode) { + this.rootNode = checkNotNull(rootNode); + return this; + } + @CanIgnoreReturnValue public Builder plugins(List plugins) { this.plugins = plugins; @@ -144,12 +173,36 @@ public App build() { if (name == null) { throw new IllegalStateException("App name must be provided."); } - if (rootAgent == null) { - throw new IllegalStateException("Root agent must be provided."); - } + BaseAgent buildRootAgent = resolveRootAgent(); validateAppName(name); return new App( - name, rootAgent, plugins, eventsCompactionConfig, contextCacheConfig, resumabilityConfig); + name, + buildRootAgent, + // An agent set as the root node is the root agent, so the app has no root node. + rootNode instanceof BaseAgent ? null : rootNode, + plugins, + eventsCompactionConfig, + contextCacheConfig, + resumabilityConfig); + } + + /** Returns the root agent, or the agent that stands in for the root node. */ + private BaseAgent resolveRootAgent() { + if (rootAgent instanceof NodeViewAgent || rootNode instanceof NodeViewAgent) { + throw new IllegalStateException( + "A node view agent cannot be the root; set the node it stands in for as the root" + + " node."); + } + if (rootNode == null) { + if (rootAgent == null) { + throw new IllegalStateException("Root agent or root node must be provided."); + } + return rootAgent; + } + if (rootAgent != null) { + throw new IllegalStateException("Set a root agent or a root node, not both."); + } + return rootNode instanceof BaseAgent agent ? agent : new NodeViewAgent(rootNode); } } 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 a7b1549b1..64f84cf32 100644 --- a/core/src/main/java/com/google/adk/runner/Runner.java +++ b/core/src/main/java/com/google/adk/runner/Runner.java @@ -27,6 +27,7 @@ import com.google.adk.agents.InvocationContext; import com.google.adk.agents.LiveRequestQueue; import com.google.adk.agents.LlmAgent; +import com.google.adk.agents.NodeViewAgent; import com.google.adk.agents.Role; import com.google.adk.agents.RunConfig; import com.google.adk.agents.SequentialAgent; @@ -55,6 +56,8 @@ import com.google.adk.tools.BaseTool; import com.google.adk.tools.FunctionTool; import com.google.adk.utils.CollectionUtils; +import com.google.adk.workflow.Node; +import com.google.adk.workflow.NodeRunner; import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableList; import com.google.common.collect.MapMaker; @@ -85,11 +88,13 @@ import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; +import java.util.function.Function; import org.jspecify.annotations.Nullable; /** The main class for the GenAI Agents runner. */ public class Runner { private final BaseAgent agent; + private final @Nullable Node node; private final String appName; private final BaseArtifactService artifactService; private final BaseSessionService sessionService; @@ -103,6 +108,8 @@ public class Runner { private final UuidProvider uuidProvider; private final ConcurrentMap activeSessionCompletables = new MapMaker().weakValues().makeMap(); + private static final String NODE_ROOT_RESUME_UNSUPPORTED = + "Resuming an invocation of a workflow node root is not supported."; /** Builder for {@link Runner}. */ public static class Builder { @@ -220,6 +227,7 @@ public Builder uuidProvider(@Nullable UuidProvider uuidProvider) { @SuppressWarnings("deprecation") public Runner build() { BaseAgent buildAgent; + Node buildNode; String buildAppName; List buildPlugins; EventsCompactionConfig buildEventsCompactionConfig; @@ -234,6 +242,7 @@ public Runner build() { throw new IllegalStateException("plugins() cannot be called when app() is called."); } buildAgent = this.app.rootAgent(); + buildNode = this.app.rootNode().orElse(null); buildPlugins = this.app.plugins(); buildAppName = this.appName == null ? this.app.name() : this.appName; buildEventsCompactionConfig = this.app.eventsCompactionConfig(); @@ -241,6 +250,7 @@ public Runner build() { buildResumabilityConfig = this.app.resumabilityConfig(); } else { buildAgent = this.agent; + buildNode = null; buildAppName = this.appName; buildPlugins = this.plugins; buildEventsCompactionConfig = null; @@ -270,6 +280,7 @@ public Runner build() { } return new Runner( buildAgent, + buildNode, buildAppName, artifactService, sessionService, @@ -364,6 +375,7 @@ protected Runner( @Nullable ResumabilityConfig resumabilityConfig) { this( agent, + /* node= */ null, appName, artifactService, sessionService, @@ -379,6 +391,7 @@ protected Runner( private Runner( BaseAgent agent, + @Nullable Node node, String appName, BaseArtifactService artifactService, BaseSessionService sessionService, @@ -390,7 +403,11 @@ private Runner( @Nullable Scheduler scheduler, InstantSource instantSource, UuidProvider uuidProvider) { + checkArgument( + !(agent instanceof NodeViewAgent) || node != null, + "A node view agent cannot run on its own; run its node through an app's root node."); this.agent = agent; + this.node = node; this.appName = appName; this.artifactService = artifactService; this.sessionService = sessionService; @@ -423,6 +440,15 @@ public BaseAgent agent() { return this.agent; } + /** + * Returns the workflow node this runner runs, or empty when it runs an agent tree. When it runs a + * node, {@link #agent()} returns a {@link NodeViewAgent} that stands in for it. + */ + @Experimental + public Optional node() { + return Optional.ofNullable(node); + } + public String appName() { return this.appName; } @@ -736,6 +762,11 @@ private Flowable runNewInvocation( @Nullable Map stateDelta) { return Flowable.defer( () -> { + // Rejected before any callback or append; defer turns the throw into an error signal. + if (node != null) { + checkNoTextWithFunctionResponses(newMessage); + checkNoAnswerToAnEarlierInvocation(session, newMessage); + } Context capturedContext = Context.current(); BaseAgent rootAgent = this.agent; String invocationId = InvocationContext.newInvocationContextId(this.uuidProvider); @@ -844,9 +875,9 @@ private Flowable runAgentWithUpdatedSession( } /** - * 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. + * Runs the root of {@code context} 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 run; {@code onEachPersisted} runs after each persisted event. */ private Flowable executeAgentPipeline( InvocationContext context, @@ -856,47 +887,59 @@ private Flowable executeAgentPipeline( // Let BaseLlmFlow block each step until this Runner has persisted the prior step's events. PersistBarrier.enable(context); - Flowable agentEvents = - context - .agent() - .runAsync(context) - .concatMap( - agentEvent -> { - // Mirror ADK Python (runners.py): partial events are streamed to the caller but - // never persisted, so managed session services (e.g. VertexAiSessionService) do - // not store a duplicate of the function call/text that the final aggregated event - // already carries. Nothing to persist, so resolve the barrier immediately. - Single persistStep = - agentEvent.partial().orElse(false) - ? Single.just(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(context, agentEvent.id())) - .doOnError( - error -> PersistBarrier.markFailed(context, agentEvent.id(), error)) - .flatMap( - registeredEvent -> { - if (onEachPersisted != null) { - onEachPersisted.run(); - } - return context - .pluginManager() - .onEventCallback(context, registeredEvent) - .defaultIfEmpty(registeredEvent); - }) - .toFlowable(); - }); - return beforeRunEvent .toFlowable() - .switchIfEmpty(agentEvents) + .switchIfEmpty( + runRoot( + context, event -> processEvent(context, sessionToPersist, onEachPersisted, event))) .concatWith(Completable.defer(() -> pluginManager.afterRunCallback(context))) .concatWith(Completable.defer(() -> compactEvents(sessionToPersist))); } + /** + * Persists {@code event} unless it is partial, runs {@code onEachPersisted} and the plugins' + * event callbacks, and returns the event to emit. + */ + private Single processEvent( + InvocationContext context, + Session sessionToPersist, + @Nullable Runnable onEachPersisted, + Event event) { + // Mirror ADK Python (runners.py): partial events are streamed to the caller but + // never persisted, so managed session services (e.g. VertexAiSessionService) do + // not store a duplicate of the function call/text that the final aggregated event + // already carries. Nothing to persist, so resolve the barrier immediately. + Single persistStep = + event.partial().orElse(false) + ? Single.just(event) + : this.sessionService.appendEvent(sessionToPersist, event); + return persistStep + // Release (or fail) BaseLlmFlow's wait for this step; the Runner stays the + // sole appendEvent caller (see PersistBarrier). + .doOnSuccess(unusedEvent -> PersistBarrier.markPersisted(context, event.id())) + .doOnError(error -> PersistBarrier.markFailed(context, event.id(), error)) + .flatMap( + registeredEvent -> { + if (onEachPersisted != null) { + onEachPersisted.run(); + } + return context + .pluginManager() + .onEventCallback(context, registeredEvent) + .defaultIfEmpty(registeredEvent); + }); + } + + /** Runs an agent root through its own lifecycle, and any other node root through NodeRunner. */ + private static Flowable runRoot( + InvocationContext context, Function> processEvent) { + Node root = context.node().orElse(context.agent()); + if (root instanceof BaseAgent agent) { + return agent.runAsync(context).concatMapSingle(processEvent::apply); + } + return NodeRunner.runRoot(root, context, processEvent); + } + private Completable compactEvents(Session session) { return Optional.ofNullable(eventsCompactionConfig) .filter(EventsCompactionConfig::hasSlidingWindowCompactionConfig) @@ -941,8 +984,14 @@ private Flowable runResumableFromSession( return Flowable.defer( () -> { // Rejected before the invocation-id fallback; defer turns the throw into an error signal. + if (node != null) { + checkNoTextWithFunctionResponses(newMessage); + } if (newMessage != null) { - validateResumeMessage(session, newMessage); + validateResumeMessage( + session, + newMessage, + /* requireFirstId= */ node == null || providedInvocationId != null); } String resolvedInvocationId = resolveInvocationId(session, newMessage, providedInvocationId); @@ -959,6 +1008,11 @@ private Flowable runResumableFromSession( return Flowable.error( new IllegalArgumentException("Session has no events to resume.")); } + if (node != null) { + // Without replay, resuming would run the node again and repeat its side effects. + return Flowable.error( + new UnsupportedOperationException(NODE_ROOT_RESUME_UNSUPPORTED)); + } return resumeCore(session, resolvedInvocationId, newMessage, runConfig, stateDelta); }); } @@ -1227,16 +1281,55 @@ private static Optional matchingFunctionCallEvent(Session session, Conten return Optional.empty(); } + /** + * Rejects a message for a node root if the message mixes text with function responses that carry + * an id: the responses resume an existing invocation, while the text starts a new one. Python's + * node path rejects such a message too; its agent path accepts it. + * + * @throws IllegalArgumentException if {@code message} mixes them. + */ + private static void checkNoTextWithFunctionResponses(@Nullable Content message) { + if (message == null || functionResponseIds(message).stream().allMatch(String::isEmpty)) { + return; + } + boolean hasText = + message.parts().stream() + .flatMap(List::stream) + .map(Part::text) + .flatMap(Optional::stream) + .anyMatch(text -> !text.isEmpty()); + checkArgument( + !hasText, + "Message cannot contain both function responses and text. Function responses resume an" + + " existing invocation while text starts a new one."); + } + + /** + * For a node root, rejects a message whose function responses answer calls from an earlier + * invocation: without replay, running the node again would repeat its side effects. As in Python + * and Kotlin, responses without an id are ignored, so they start a new invocation. + * + * @throws IllegalArgumentException if a response id matches no call, or the ids span invocations. + * @throws UnsupportedOperationException if the responses answer an earlier invocation. + */ + private static void checkNoAnswerToAnEarlierInvocation(Session session, Content message) { + validateResumeMessage(session, message, /* requireFirstId= */ false); + if (matchingFunctionCallEvent(session, message).isPresent()) { + throw new UnsupportedOperationException(NODE_ROOT_RESUME_UNSUPPORTED); + } + } + /** * 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. + * carries no id when {@code requireFirstId} is set, 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) { + private static void validateResumeMessage( + Session session, Content newMessage, boolean requireFirstId) { ImmutableList responses = newMessage.parts().stream() .flatMap(List::stream) @@ -1249,7 +1342,7 @@ private static void validateResumeMessage(Session session, Content newMessage) { // 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(), + !requireFirstId || responses.get(0).id().isPresent(), "A function response id is required to resume an invocation."); Set unmatched = functionResponseIds(newMessage); Set invocationIds = new HashSet<>(); @@ -1380,7 +1473,8 @@ private InvocationContext.Builder newInvocationContextBuilder(Session session, B .scheduler(this.scheduler) .instantSource(this.instantSource) .uuidProvider(this.uuidProvider) - .agent(agent); + .agent(agent) + .node(node); } public Flowable runLive( @@ -1437,6 +1531,11 @@ protected Flowable runLiveImpl( Session session, @Nullable LiveRequestQueue liveRequestQueue, RunConfig runConfig) { return Flowable.defer( () -> { + if (node != null) { + return Flowable.error( + new UnsupportedOperationException( + "Live runs of a workflow node root are not supported.")); + } Context capturedContext = Context.current(); InvocationContext invocationContext = newInvocationContextForLive(session, liveRequestQueue, runConfig); @@ -1511,6 +1610,10 @@ private boolean isLegacyResumability() { /** Returns the agent that should handle the next request based on session history. */ private BaseAgent findAgentToRun(Session session, BaseAgent rootAgent) { + if (node != null) { + // A node root has no agent tree to pick a sub-agent from, so it always runs itself. + return rootAgent; + } // 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. ImmutableList sessionEvents = session.immutableEvents(); diff --git a/core/src/main/java/com/google/adk/workflow/NodeRunner.java b/core/src/main/java/com/google/adk/workflow/NodeRunner.java index bd2284a56..4c7a0a31b 100644 --- a/core/src/main/java/com/google/adk/workflow/NodeRunner.java +++ b/core/src/main/java/com/google/adk/workflow/NodeRunner.java @@ -32,16 +32,23 @@ import io.reactivex.rxjava3.core.Completable; import io.reactivex.rxjava3.core.Flowable; import io.reactivex.rxjava3.core.Single; +import io.reactivex.rxjava3.processors.FlowableProcessor; +import io.reactivex.rxjava3.processors.UnicastProcessor; +import io.reactivex.rxjava3.subjects.CompletableSubject; import java.util.Map; +import java.util.Queue; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.function.Function; import org.jspecify.annotations.Nullable; /** * The runner for one workflow node. It stamps every emitted event with its author, node path, and * branch before forwarding it, and records the node's output, routes, interrupts, and failure on - * the node state of the returned context. + * the node state of the returned context. {@link #runRoot} runs a node as the root of an + * invocation. */ @Experimental -final class NodeRunner { +public final class NodeRunner { /** Sink that receives the events of a workflow run. */ @Experimental @@ -49,10 +56,11 @@ interface EventSink { /** * Sends {@code event} and returns a {@link Completable} that completes once the event is - * processed. The Completable must not fail; a consumer that cannot process an event stops the - * run by canceling it. + * processed. When {@code nodeFailure} is not null, {@code event} is the error event of a node + * that failed with {@code nodeFailure}. The Completable must not fail; a consumer that cannot + * process an event stops the run by canceling it. */ - Completable send(Event event); + Completable send(Event event, @Nullable Exception nodeFailure); } private final Node node; @@ -90,6 +98,77 @@ static Builder builder() { return new Builder(); } + /** + * Runs {@code node} as the root of {@code context} on the invocation's user content, passing each + * event through {@code processEvent} in order; only {@link com.google.adk.runner.Runner} should + * call this method. The node is asked for its next event, and its failure is signaled, only after + * the previous non-partial event is processed and emitted, so a failure follows the error event + * that reports it. If processing an event fails with an exception, the failures reported by + * unprocessed error events are attached to it as suppressed. + */ + @Experimental + public static Flowable runRoot( + Node node, + InvocationContext context, + Function> processEvent) { + checkNotNull(node); + checkNotNull(context); + checkNotNull(processEvent); + return Flowable.defer( + () -> { + FlowableProcessor queue = + UnicastProcessor.create().toSerialized(); + Queue unprocessedFailures = new ConcurrentLinkedQueue<>(); + EventSink sink = + (event, nodeFailure) -> { + CompletableSubject processed = CompletableSubject.create(); + if (nodeFailure != null) { + unprocessedFailures.add(nodeFailure); + } + queue.onNext(new QueuedEvent(event, processed, nodeFailure)); + return event.partial().orElse(false) ? Completable.complete() : processed; + }; + Completable nodeRun = + builder() + .node(node) + .invocationContext(context) + .eventSink(sink) + .build() + .run(context.userContent().orElse(null)) + .doOnSuccess( + nodeContext -> + nodeContext + .requireNodeState() + .failure() + .ifPresentOrElse(queue::onError, queue::onComplete)) + .doOnError(queue::onError) + .ignoreElement() + .onErrorComplete(); + return queue + .concatMap( + queued -> + processEvent + .apply(queued.event()) + .doOnError( + error -> { + // As in Kotlin, keep the failures of unprocessed error events. + if (error instanceof Exception) { + unprocessedFailures.forEach(error::addSuppressed); + } + }) + .toFlowable() + // The node proceeds only after the caller receives the processed event. + .doAfterNext( + unused -> { + if (queued.nodeFailure() != null) { + unprocessedFailures.remove(queued.nodeFailure()); + } + queued.processed().onComplete(); + })) + .mergeWith(nodeRun); + }); + } + /** * Runs the node on {@code nodeInput} and returns its context when the node finishes. An exception * from the node is recorded on the context's node state rather than signaled, while an {@code @@ -190,7 +269,7 @@ private Completable send(Event event) { Event stamped = stamp(messageAsOutput ? event.toBuilder().output(null).build() : event); Event outgoing = stamped.partial().orElse(false) ? stamped : withPendingDeltas(stamped); return eventSink - .send(outgoing) + .send(outgoing, /* nodeFailure= */ null) .doOnComplete( () -> { if (outgoing.output().isPresent() || messageAsOutput) { @@ -219,7 +298,7 @@ private Completable flushPending() { .build()) .build())); return eventSink - .send(event) + .send(event, /* nodeFailure= */ null) .doOnComplete( () -> { if (pendingOutput) { @@ -243,7 +322,7 @@ private Completable fail(Exception failure) { .errorCode(new FinishReason(errorTypeName(failure))) .errorMessage(nullToEmpty(failure.getMessage())) .build())); - return eventSink.send(errorEvent).doOnComplete(() -> nodeState.failure(failure)); + return eventSink.send(errorEvent, failure).doOnComplete(() -> nodeState.failure(failure)); } /** @@ -290,6 +369,13 @@ private Event.Builder newEvent() { } } + /** + * An event queued for the runner, paired with the signal that completes once the caller receives + * the event and, for an error event, with the failure it reports. + */ + private record QueuedEvent( + Event event, CompletableSubject processed, @Nullable Exception nodeFailure) {} + /** Appends a {@code name@runId} segment to {@code base}, joined by {@code separator}. */ static String appendSegment(@Nullable String base, String name, String runId, char separator) { String segment = name + "@" + runId; diff --git a/core/src/test/java/com/google/adk/agents/BaseAgentTest.java b/core/src/test/java/com/google/adk/agents/BaseAgentTest.java index 653749023..4d803edbe 100644 --- a/core/src/test/java/com/google/adk/agents/BaseAgentTest.java +++ b/core/src/test/java/com/google/adk/agents/BaseAgentTest.java @@ -25,7 +25,9 @@ import com.google.adk.telemetry.Metrics; import com.google.adk.testing.TestBaseAgent; import com.google.adk.testing.TestCallback; +import com.google.adk.testing.TestNode; import com.google.adk.testing.TestUtils; +import com.google.adk.workflow.Node; import com.google.common.collect.ImmutableList; import com.google.genai.types.Content; import com.google.genai.types.Part; @@ -406,6 +408,21 @@ public void runAsync_firstAfterCallbackReturnsContent_skipsSecondAfterCallback() assertThat(afterCallback2.wasCalled()).isFalse(); } + @Test + public void runAsync_underAnInvocationOfANode_runsWithoutTheNode() { + TestBaseAgent agent = TestUtils.createSubAgent(TEST_AGENT_NAME); + Node node = TestNode.of("root"); + InvocationContext parentContext = + TestUtils.createInvocationContext(agent).toBuilder() + .agent(new NodeViewAgent(node)) + .node(node) + .build(); + + agent.runAsync(parentContext).blockingSubscribe(); + + assertThat(agent.getLastInvocationContext().node()).isEmpty(); + } + @Test public void canonicalCallbacks_returnsEmptyListWhenNull() { TestBaseAgent agent = 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 8e19ce355..b9232a976 100644 --- a/core/src/test/java/com/google/adk/agents/InvocationContextTest.java +++ b/core/src/test/java/com/google/adk/agents/InvocationContextTest.java @@ -31,6 +31,8 @@ import com.google.adk.sessions.BaseSessionService; import com.google.adk.sessions.Session; import com.google.adk.summarizer.EventsCompactionConfig; +import com.google.adk.testing.TestNode; +import com.google.adk.workflow.Node; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; @@ -106,6 +108,53 @@ public void isResumable_andIsLegacyResumability_separateTheTwoModes() { assertThat(neitherContext.isLegacyResumability()).isFalse(); } + @Test + public void node_isTheNodeSetOnTheBuilder() { + Node node = TestNode.of("solo"); + + InvocationContext context = + InvocationContext.builder() + .sessionService(mockSessionService) + .agent(new NodeViewAgent(node)) + .node(node) + .session(session) + .build(); + + assertThat(context.node()).hasValue(node); + } + + @Test + public void toBuilder_keepsTheNode() { + Node node = TestNode.of("solo"); + InvocationContext context = + InvocationContext.builder() + .sessionService(mockSessionService) + .agent(new NodeViewAgent(node)) + .node(node) + .session(session) + .build(); + + assertThat(context.toBuilder().build().node()).hasValue(node); + } + + @Test + public void build_withANodeViewAgentButNoNode_throwsIllegalStateException() { + InvocationContext.Builder builder = + InvocationContext.builder() + .sessionService(mockSessionService) + .agent(new NodeViewAgent(TestNode.of("solo"))) + .session(session); + + assertThrows(IllegalStateException.class, builder::build); + } + + @Test + public void node_whenTheAgentIsAnAgent_isEmpty() { + InvocationContext context = contextWith(/* resumabilityConfig= */ null); + + assertThat(context.node()).isEmpty(); + } + private InvocationContext contextWith(ResumabilityConfig resumabilityConfig) { return InvocationContext.builder() .sessionService(mockSessionService) diff --git a/core/src/test/java/com/google/adk/agents/NodeViewAgentTest.java b/core/src/test/java/com/google/adk/agents/NodeViewAgentTest.java new file mode 100644 index 000000000..18c704330 --- /dev/null +++ b/core/src/test/java/com/google/adk/agents/NodeViewAgentTest.java @@ -0,0 +1,92 @@ +/* + * 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.agents; + +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; + +import com.google.adk.testing.TestNode; +import com.google.adk.testing.TestUtils; +import com.google.adk.workflow.Node; +import io.reactivex.rxjava3.core.Flowable; +import org.jspecify.annotations.Nullable; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public final class NodeViewAgentTest { + + private static final Node NODE = + new Node() { + @Override + public String name() { + return "solo"; + } + + @Override + public String description() { + return "A test node."; + } + + @Override + public Flowable runNode(Context context, @Nullable Object nodeInput) { + return Flowable.empty(); + } + }; + + @Test + public void constructor_takesTheNameAndDescriptionOfTheNode() { + NodeViewAgent agent = new NodeViewAgent(NODE); + + assertThat(agent.name()).isEqualTo("solo"); + assertThat(agent.description()).isEqualTo("A test node."); + assertThat(agent.node()).isSameInstanceAs(NODE); + } + + @Test + public void constructor_withAnAgent_throwsIllegalArgumentException() { + BaseAgent agent = TestUtils.createSubAgent("worker"); + + assertThrows(IllegalArgumentException.class, () -> new NodeViewAgent(agent)); + } + + @Test + public void constructor_withANodeNameThatIsNotAValidAgentName_throwsIllegalArgumentException() { + Node node = TestNode.of("not@valid"); + + assertThrows(IllegalArgumentException.class, () -> new NodeViewAgent(node)); + } + + @Test + public void runAsync_failsWithUnsupportedOperationException() { + NodeViewAgent agent = new NodeViewAgent(NODE); + + agent.runAsync(parentContext()).test().assertError(UnsupportedOperationException.class); + } + + @Test + public void runLive_failsWithUnsupportedOperationException() { + NodeViewAgent agent = new NodeViewAgent(NODE); + + agent.runLive(parentContext()).test().assertError(UnsupportedOperationException.class); + } + + private static InvocationContext parentContext() { + return TestUtils.createInvocationContext(TestUtils.createRootAgent()); + } +} diff --git a/core/src/test/java/com/google/adk/apps/AppTest.java b/core/src/test/java/com/google/adk/apps/AppTest.java new file mode 100644 index 000000000..7202e2233 --- /dev/null +++ b/core/src/test/java/com/google/adk/apps/AppTest.java @@ -0,0 +1,98 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.adk.apps; + +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; + +import com.google.adk.agents.BaseAgent; +import com.google.adk.agents.NodeViewAgent; +import com.google.adk.testing.TestNode; +import com.google.adk.testing.TestUtils; +import com.google.adk.workflow.Node; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public final class AppTest { + + private static final Node NODE = TestNode.of("solo"); + + @Test + public void build_withARootNode_standsANodeViewAgentInForIt() { + App app = App.builder().name("node_app").rootNode(NODE).build(); + + assertThat(app.rootNode()).hasValue(NODE); + assertThat(app.rootAgent()).isInstanceOf(NodeViewAgent.class); + assertThat(((NodeViewAgent) app.rootAgent()).node()).isSameInstanceAs(NODE); + } + + @Test + public void build_withAnAgentAsTheRootNode_makesItTheRootAgent() { + BaseAgent agent = TestUtils.createSubAgent("worker"); + + App app = App.builder().name("agent_app").rootNode(agent).build(); + + assertThat(app.rootAgent()).isSameInstanceAs(agent); + assertThat(app.rootNode()).isEmpty(); + } + + @Test + public void build_withARootAgent_hasNoRootNode() { + App app = App.builder().name("agent_app").rootAgent(TestUtils.createSubAgent("worker")).build(); + + assertThat(app.rootNode()).isEmpty(); + } + + @Test + public void build_withARootAgentAndADifferentRootNode_throwsIllegalStateException() { + App.Builder builder = + App.builder().name("app").rootAgent(TestUtils.createSubAgent("worker")).rootNode(NODE); + + assertThrows(IllegalStateException.class, builder::build); + } + + @Test + public void build_withTheSameAgentAsRootAgentAndRootNode_throwsIllegalStateException() { + BaseAgent agent = TestUtils.createSubAgent("worker"); + App.Builder builder = App.builder().name("app").rootAgent(agent).rootNode(agent); + + assertThrows(IllegalStateException.class, builder::build); + } + + @Test + public void build_withNeitherARootAgentNorARootNode_throwsIllegalStateException() { + App.Builder builder = App.builder().name("app"); + + assertThrows(IllegalStateException.class, builder::build); + } + + @Test + public void build_withANodeViewAgentAsTheRootAgent_throwsIllegalStateException() { + App.Builder builder = App.builder().name("app").rootAgent(new NodeViewAgent(NODE)); + + assertThrows(IllegalStateException.class, builder::build); + } + + @Test + public void build_withANodeViewAgentAsTheRootNode_throwsIllegalStateException() { + App.Builder builder = App.builder().name("app").rootNode(new NodeViewAgent(NODE)); + + assertThrows(IllegalStateException.class, builder::build); + } +} diff --git a/core/src/test/java/com/google/adk/runner/RunnerRootNodeTest.java b/core/src/test/java/com/google/adk/runner/RunnerRootNodeTest.java new file mode 100644 index 000000000..a931e13ef --- /dev/null +++ b/core/src/test/java/com/google/adk/runner/RunnerRootNodeTest.java @@ -0,0 +1,628 @@ +/* + * 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.runner; + +import static com.google.common.collect.ImmutableList.toImmutableList; +import static com.google.common.collect.Iterables.getLast; +import static com.google.common.collect.Iterables.getOnlyElement; +import static com.google.common.truth.Truth.assertThat; +import static java.util.concurrent.TimeUnit.SECONDS; +import static org.junit.Assert.assertThrows; + +import com.google.adk.agents.BaseAgent; +import com.google.adk.agents.CallbackContext; +import com.google.adk.agents.InvocationContext; +import com.google.adk.agents.LiveRequestQueue; +import com.google.adk.agents.NodeViewAgent; +import com.google.adk.agents.RunConfig; +import com.google.adk.apps.App; +import com.google.adk.apps.ResumabilityConfig; +import com.google.adk.artifacts.InMemoryArtifactService; +import com.google.adk.events.Event; +import com.google.adk.plugins.BasePlugin; +import com.google.adk.plugins.Plugin; +import com.google.adk.sessions.InMemorySessionService; +import com.google.adk.sessions.Session; +import com.google.adk.testing.TestBaseAgent; +import com.google.adk.testing.TestNode; +import com.google.adk.workflow.Node; +import com.google.adk.workflow.NodeInfo; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.genai.types.Content; +import com.google.genai.types.FinishReason; +import com.google.genai.types.FunctionCall; +import com.google.genai.types.FunctionResponse; +import com.google.genai.types.Part; +import io.reactivex.rxjava3.core.Completable; +import io.reactivex.rxjava3.core.Flowable; +import io.reactivex.rxjava3.core.Maybe; +import io.reactivex.rxjava3.subscribers.TestSubscriber; +import java.util.ArrayList; +import java.util.List; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests for running a workflow node, rather than an agent, at the root of an invocation. */ +@RunWith(JUnit4.class) +public final class RunnerRootNodeTest { + + private static final String APP_NAME = "node_app"; + private static final String USER_ID = "user"; + private static final Content USER_MESSAGE = Content.fromParts(Part.fromText("hello")); + private static final FunctionResponse IDENTIFIED_RESPONSE = + FunctionResponse.builder().id("fc-1").name("tool").response(ImmutableMap.of("v", 1)).build(); + private static final FunctionResponse UNIDENTIFIED_RESPONSE = + FunctionResponse.builder().name("tool").response(ImmutableMap.of("v", 1)).build(); + + @Test + public void node_ofARunnerForANodeApp_isThatNode() { + Node node = TestNode.of("solo"); + + Runner runner = runner(App.builder().name(APP_NAME).rootNode(node).build()); + + assertThat(runner.node()).hasValue(node); + } + + @Test + public void build_withANodeViewAgentAsTheAgent_throwsIllegalArgumentException() { + Runner.Builder builder = + Runner.builder().agent(new NodeViewAgent(TestNode.of("solo"))).appName(APP_NAME); + + assertThrows(IllegalArgumentException.class, builder::build); + } + + @Test + @SuppressWarnings("deprecation") // The deprecated constructors must reject it as well. + public void constructor_withANodeViewAgent_throwsIllegalArgumentException() { + NodeViewAgent agent = new NodeViewAgent(TestNode.of("solo")); + + assertThrows( + IllegalArgumentException.class, + () -> + new Runner( + agent, APP_NAME, new InMemoryArtifactService(), new InMemorySessionService())); + } + + @Test + public void runAsync_nodeRoot_runsTheNodeOnTheUserMessage() { + List inputs = new ArrayList<>(); + Node node = + new TestNode( + "solo", + (context, input) -> { + inputs.add(input); + return Flowable.empty(); + }); + Runner runner = runner(node); + Session session = newSession(runner); + + runner.runAsync(USER_ID, session.id(), USER_MESSAGE).blockingSubscribe(); + + assertThat(inputs).containsExactly(USER_MESSAGE); + } + + @Test + public void runAsync_nodeRoot_persistsItsStampedEventAndState() { + Node node = + new TestNode( + "solo", + (context, input) -> { + context.state().put("durable", "saved"); + return Flowable.just("echoed"); + }); + Runner runner = runner(node); + Session session = newSession(runner); + + List events = + runner.runAsync(USER_ID, session.id(), USER_MESSAGE).toList().blockingGet(); + + Event event = getOnlyElement(events); + Session saved = reload(runner, session); + assertThat(event.author()).isEqualTo("solo"); + assertThat(event.nodeInfo().map(NodeInfo::path)).hasValue("solo@1"); + assertThat(event.output()).hasValue("echoed"); + assertThat(getLast(saved.immutableEvents()).id()).isEqualTo(event.id()); + assertThat(saved.state()).containsEntry("durable", "saved"); + } + + @Test + public void runAsync_failingNodeRoot_persistsItsErrorEventBeforeFailing() { + IllegalStateException failure = new IllegalStateException("node failed"); + Runner runner = runner(failingNode(failure)); + Session session = newSession(runner); + + TestSubscriber subscriber = runner.runAsync(USER_ID, session.id(), USER_MESSAGE).test(); + subscriber.awaitDone(5, SECONDS); + + subscriber.assertError(failure); + Event errorEvent = getOnlyElement(subscriber.values()); + Session saved = reload(runner, session); + assertThat(errorEvent.errorCode()).hasValue(new FinishReason("IllegalStateException")); + assertThat(errorEvent.nodeInfo().map(NodeInfo::path)).hasValue("boom@1"); + assertThat(getLast(saved.immutableEvents()).id()).isEqualTo(errorEvent.id()); + assertThat(saved.state()).containsEntry("partial_step", "done"); + } + + @Test + public void runAsync_failingNodeRoot_reportsTheErrorToPluginsWithoutAfterRun() { + IllegalStateException failure = new IllegalStateException("node failed"); + RecordingPlugin plugin = new RecordingPlugin(); + Runner runner = runner(failingNode(failure), plugin); + Session session = newSession(runner); + + TestSubscriber subscriber = runner.runAsync(USER_ID, session.id(), USER_MESSAGE).test(); + subscriber.awaitDone(5, SECONDS); + + Event errorEvent = getOnlyElement(subscriber.values()); + assertThat(ids(plugin.events)).containsExactly(errorEvent.id()); + assertThat(plugin.runErrors).containsExactly(failure); + assertThat(plugin.afterRunCount).isEqualTo(0); + } + + @Test + public void runAsync_onEventCallbackFailure_cancelsTheNodeRootAndFails() { + IllegalStateException callbackFailure = new IllegalStateException("onEvent failed"); + AtomicBoolean cancelled = new AtomicBoolean(); + Node node = + new TestNode( + "emitter", + (context, input) -> + Flowable.just(contentEvent("first")) + .concatWith(Flowable.never()) + .doOnCancel(() -> cancelled.set(true))); + Plugin failingPlugin = + new BasePlugin("failing") { + @Override + public Maybe onEventCallback(InvocationContext invocationContext, Event event) { + return Maybe.error(callbackFailure); + } + }; + Runner runner = runner(node, failingPlugin); + Session session = newSession(runner); + + TestSubscriber subscriber = runner.runAsync(USER_ID, session.id(), USER_MESSAGE).test(); + subscriber.awaitDone(5, SECONDS); + + subscriber.assertError(callbackFailure); + subscriber.assertNoValues(); + assertThat(cancelled.get()).isTrue(); + } + + @Test + public void runAsync_beforeRunCallbackContent_skipsTheNodeRootAndRunsAfterRun() { + Content halted = Content.fromParts(Part.fromText("halted by plugin")); + AtomicBoolean nodeRan = new AtomicBoolean(); + Node node = + new TestNode( + "should_not_run", + (context, input) -> { + nodeRan.set(true); + return Flowable.just("unreachable"); + }); + RecordingPlugin plugin = + new RecordingPlugin() { + @Override + public Maybe beforeRunCallback(InvocationContext invocationContext) { + return Maybe.just(halted); + } + }; + Runner runner = runner(node, plugin); + Session session = newSession(runner); + + List events = + runner.runAsync(USER_ID, session.id(), USER_MESSAGE).toList().blockingGet(); + + assertThat(nodeRan.get()).isFalse(); + assertThat(getOnlyElement(events).content()).hasValue(halted); + assertThat(plugin.afterRunCount).isEqualTo(1); + } + + @Test + public void runAsync_onUserMessageCallback_replacesTheNodeRootsInput() { + Content rewritten = Content.fromParts(Part.fromText("modified text")); + List inputs = new ArrayList<>(); + Node node = + new TestNode( + "recorder", + (context, input) -> { + inputs.add(input); + return Flowable.empty(); + }); + Plugin rewritingPlugin = + new BasePlugin("rewriting") { + @Override + public Maybe onUserMessageCallback( + InvocationContext invocationContext, Content userMessage) { + return Maybe.just(rewritten); + } + }; + Runner runner = runner(node, rewritingPlugin); + Session session = newSession(runner); + + runner.runAsync(USER_ID, session.id(), USER_MESSAGE).blockingSubscribe(); + + assertThat(inputs).containsExactly(rewritten); + } + + @Test + public void runAsync_partialEventOfANodeRoot_isStreamedButNotPersisted() { + Node node = + new TestNode( + "streamer", + (context, input) -> + Flowable.just( + contentEvent("part").toBuilder().partial(true).build(), contentEvent("final"))); + Runner runner = runner(node); + Session session = newSession(runner); + + List events = + runner.runAsync(USER_ID, session.id(), USER_MESSAGE).toList().blockingGet(); + + assertThat(texts(events)).containsExactly("part", "final").inOrder(); + assertThat( + texts( + reload(runner, session).immutableEvents().stream() + .filter(event -> event.author().equals("streamer")) + .collect(toImmutableList()))) + .containsExactly("final"); + } + + @Test + public void runAsync_takingTheFirstEventOfANodeRoot_cancelsTheNode() { + AtomicBoolean cancelled = new AtomicBoolean(); + Node node = + new TestNode( + "streamer", + (context, input) -> + Flowable.just(contentEvent("first")) + .concatWith(Flowable.never()) + .doOnCancel(() -> cancelled.set(true))); + Runner runner = runner(node); + Session session = newSession(runner); + + TestSubscriber subscriber = + runner.runAsync(USER_ID, session.id(), USER_MESSAGE).take(1).test(); + subscriber.awaitDone(5, SECONDS); + + subscriber.assertComplete(); + assertThat(texts(subscriber.values())).containsExactly("first"); + assertThat(cancelled.get()).isTrue(); + } + + @Test + public void runAsync_resumingANodeRootInvocation_failsWithoutRunningTheNodeAgain() { + AtomicInteger runs = new AtomicInteger(); + Node node = + new TestNode( + "solo", + (context, input) -> { + runs.incrementAndGet(); + return Flowable.just("done"); + }); + Runner runner = resumableRunner(node); + Session session = newSession(runner); + String invocationId = + runner.runAsync(USER_ID, session.id(), USER_MESSAGE).blockingFirst().invocationId(); + + TestSubscriber subscriber = + runner + .runAsync( + USER_ID, + session.id(), + invocationId, + /* newMessage= */ null, + RunConfig.builder().build(), + /* stateDelta= */ null) + .test(); + subscriber.awaitDone(5, SECONDS); + + subscriber.assertError(UnsupportedOperationException.class); + assertThat(runs.get()).isEqualTo(1); + } + + @Test + public void runAsync_textWithAnIdentifiedFunctionResponseToANodeRoot_failsBeforeAnyEvent() { + AtomicInteger runs = new AtomicInteger(); + Runner runner = runner(countingNode(runs)); + Session session = newSession(runner); + Content message = textWithFunctionResponse(IDENTIFIED_RESPONSE); + Flowable events = runner.runAsync(USER_ID, session.id(), message); + + IllegalArgumentException error = + assertThrows(IllegalArgumentException.class, events::blockingSubscribe); + + assertThat(error).hasMessageThat().contains("both function responses and text"); + assertThat(runs.get()).isEqualTo(0); + assertThat(reload(runner, session).immutableEvents()).isEmpty(); + } + + @Test + public void runAsync_textWithAnIdentifiedFunctionResponseToAResumableNodeRoot_fails() { + AtomicInteger runs = new AtomicInteger(); + Runner runner = resumableRunner(countingNode(runs)); + Session session = newSession(runner); + Content message = textWithFunctionResponse(IDENTIFIED_RESPONSE); + Flowable events = runner.runAsync(USER_ID, session.id(), message); + + IllegalArgumentException error = + assertThrows(IllegalArgumentException.class, events::blockingSubscribe); + + assertThat(error).hasMessageThat().contains("both function responses and text"); + assertThat(runs.get()).isEqualTo(0); + } + + @Test + public void runAsync_textWithAFunctionResponseWithoutIdToANodeRoot_runsTheNode() { + AtomicInteger runs = new AtomicInteger(); + Runner runner = runner(countingNode(runs)); + Session session = newSession(runner); + + runner + .runAsync(USER_ID, session.id(), textWithFunctionResponse(UNIDENTIFIED_RESPONSE)) + .blockingSubscribe(); + + assertThat(runs.get()).isEqualTo(1); + } + + @Test + public void runAsync_functionResponseToTheNodeRootsEarlierCall_failsWithoutRunningTheNodeAgain() { + AtomicInteger runs = new AtomicInteger(); + Node node = + new TestNode( + "solo", + (context, input) -> { + runs.incrementAndGet(); + return Flowable.just(functionCallEvent(IDENTIFIED_RESPONSE.id().orElseThrow())); + }); + Runner runner = runner(node); + Session session = newSession(runner); + runner.runAsync(USER_ID, session.id(), USER_MESSAGE).blockingSubscribe(); + int eventsAfterTheFirstRun = reload(runner, session).immutableEvents().size(); + Flowable events = + runner.runAsync(USER_ID, session.id(), functionResponse(IDENTIFIED_RESPONSE)); + + assertThrows(UnsupportedOperationException.class, events::blockingSubscribe); + + assertThat(runs.get()).isEqualTo(1); + assertThat(reload(runner, session).immutableEvents()).hasSize(eventsAfterTheFirstRun); + } + + @Test + public void runAsync_functionResponseWithAnUnknownIdToANodeRoot_failsBeforeAnyEvent() { + AtomicInteger runs = new AtomicInteger(); + Runner runner = runner(countingNode(runs)); + Session session = newSession(runner); + Flowable events = + runner.runAsync(USER_ID, session.id(), functionResponse(IDENTIFIED_RESPONSE)); + + IllegalArgumentException error = + assertThrows(IllegalArgumentException.class, events::blockingSubscribe); + + assertThat(error).hasMessageThat().contains("No matching function call"); + assertThat(runs.get()).isEqualTo(0); + assertThat(reload(runner, session).immutableEvents()).isEmpty(); + } + + @Test + public void runAsync_functionResponseWithoutIdToAResumableNodeRoot_runsTheNode() { + AtomicInteger runs = new AtomicInteger(); + Runner runner = resumableRunner(countingNode(runs)); + Session session = newSession(runner); + + runner + .runAsync(USER_ID, session.id(), functionResponse(UNIDENTIFIED_RESPONSE)) + .blockingSubscribe(); + + assertThat(runs.get()).isEqualTo(1); + } + + @Test + public void runAsync_textWithAFunctionResponseWithoutIdToAResumableNodeRoot_runsTheNode() { + AtomicInteger runs = new AtomicInteger(); + Runner runner = resumableRunner(countingNode(runs)); + Session session = newSession(runner); + + runner + .runAsync(USER_ID, session.id(), textWithFunctionResponse(UNIDENTIFIED_RESPONSE)) + .blockingSubscribe(); + + assertThat(runs.get()).isEqualTo(1); + } + + @Test + public void runAsync_resumingANodeRootWithAFunctionResponseWithoutId_failsBeforeRunningTheNode() { + AtomicInteger runs = new AtomicInteger(); + Runner runner = resumableRunner(countingNode(runs)); + Session session = newSession(runner); + runner.runAsync(USER_ID, session.id(), USER_MESSAGE).blockingSubscribe(); + String invocationId = getLast(reload(runner, session).immutableEvents()).invocationId(); + Flowable events = + runner.runAsync( + USER_ID, + session.id(), + invocationId, + functionResponse(UNIDENTIFIED_RESPONSE), + RunConfig.builder().build(), + /* stateDelta= */ null); + + IllegalArgumentException error = + assertThrows(IllegalArgumentException.class, events::blockingSubscribe); + + assertThat(error).hasMessageThat().contains("function response id is required"); + assertThat(runs.get()).isEqualTo(1); + } + + @Test + public void runAsync_textWithAnIdentifiedFunctionResponseToAnAgentRoot_runsTheAgent() { + TestBaseAgent agent = new TestBaseAgent("solo_agent", "An agent.", null, null, Flowable::empty); + Runner runner = runner(App.builder().name(APP_NAME).rootAgent(agent).build()); + Session session = newSession(runner); + + runner + .runAsync(USER_ID, session.id(), textWithFunctionResponse(IDENTIFIED_RESPONSE)) + .blockingSubscribe(); + + assertThat(agent.getInvocationCount()).isEqualTo(1); + } + + @Test + public void runLive_nodeRoot_failsWithUnsupportedOperationExceptionBeforeAnyAgentCallback() { + AtomicBoolean agentCallbackRan = new AtomicBoolean(); + Plugin agentCallbackRecorder = + new BasePlugin("agent_callback_recorder") { + @Override + public Maybe beforeAgentCallback( + BaseAgent agent, CallbackContext callbackContext) { + agentCallbackRan.set(true); + return Maybe.empty(); + } + }; + Runner runner = runner(TestNode.of("solo"), agentCallbackRecorder); + Session session = newSession(runner); + + TestSubscriber subscriber = + runner.runLive(session, new LiveRequestQueue(), RunConfig.builder().build()).test(); + subscriber.awaitDone(5, SECONDS); + + subscriber.assertError(UnsupportedOperationException.class); + assertThat(agentCallbackRan.get()).isFalse(); + } + + private static Runner runner(Node node, Plugin... plugins) { + return runner(App.builder().name(APP_NAME).rootNode(node).plugins(plugins).build()); + } + + private static Runner runner(App app) { + return Runner.builder().app(app).build(); + } + + private static Runner resumableRunner(Node node) { + return runner( + App.builder() + .name(APP_NAME) + .rootNode(node) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build()); + } + + private static Session newSession(Runner runner) { + return runner.sessionService().createSession(APP_NAME, USER_ID).blockingGet(); + } + + private static Session reload(Runner runner, Session session) { + return runner + .sessionService() + .getSession(APP_NAME, USER_ID, session.id(), Optional.empty()) + .blockingGet(); + } + + private static Node failingNode(Exception failure) { + return new TestNode( + "boom", + (context, input) -> { + context.state().put("partial_step", "done"); + return Flowable.error(failure); + }); + } + + private static Node countingNode(AtomicInteger runs) { + return new TestNode( + "solo", + (context, input) -> { + runs.incrementAndGet(); + return Flowable.empty(); + }); + } + + private static Content functionResponse(FunctionResponse response) { + return Content.builder() + .role("user") + .parts(Part.builder().functionResponse(response).build()) + .build(); + } + + private static Content textWithFunctionResponse(FunctionResponse response) { + return Content.builder() + .role("user") + .parts(Part.fromText("some text"), Part.builder().functionResponse(response).build()) + .build(); + } + + private static Event contentEvent(String text) { + return Event.builder() + .id(Event.generateEventId()) + .author("") + .content(Content.fromParts(Part.fromText(text))) + .build(); + } + + private static Event functionCallEvent(String callId) { + return Event.builder() + .id(Event.generateEventId()) + .author("") + .content( + Content.fromParts( + Part.builder() + .functionCall(FunctionCall.builder().id(callId).name("tool").build()) + .build())) + .build(); + } + + private static ImmutableList ids(List events) { + return events.stream().map(Event::id).collect(toImmutableList()); + } + + private static ImmutableList texts(List events) { + return events.stream().map(Event::stringifyContent).collect(toImmutableList()); + } + + /** + * A plugin that records the events, run errors, and completed runs reported to it by the runner. + */ + private static class RecordingPlugin extends BasePlugin { + final List events = new ArrayList<>(); + final List runErrors = new ArrayList<>(); + int afterRunCount; + + RecordingPlugin() { + super("recording"); + } + + @Override + public Maybe onEventCallback(InvocationContext invocationContext, Event event) { + events.add(event); + return Maybe.empty(); + } + + @Override + public Completable onRunErrorCallback(InvocationContext invocationContext, Throwable error) { + runErrors.add(error); + return Completable.complete(); + } + + @Override + public Completable afterRunCallback(InvocationContext invocationContext) { + afterRunCount++; + return Completable.complete(); + } + } +} diff --git a/core/src/test/java/com/google/adk/workflow/NodeExecutionStateTest.java b/core/src/test/java/com/google/adk/workflow/NodeExecutionStateTest.java index 9ad3e60c9..1463c973a 100644 --- a/core/src/test/java/com/google/adk/workflow/NodeExecutionStateTest.java +++ b/core/src/test/java/com/google/adk/workflow/NodeExecutionStateTest.java @@ -27,7 +27,7 @@ @RunWith(JUnit4.class) public final class NodeExecutionStateTest { - private static final NodeRunner.EventSink SINK = event -> Completable.complete(); + private static final NodeRunner.EventSink SINK = (event, nodeFailure) -> Completable.complete(); @Test public void build_withoutANode_throwsIllegalStateException() { diff --git a/core/src/test/java/com/google/adk/workflow/NodeRunnerTest.java b/core/src/test/java/com/google/adk/workflow/NodeRunnerTest.java index 5f04c5a07..3e472dc5e 100644 --- a/core/src/test/java/com/google/adk/workflow/NodeRunnerTest.java +++ b/core/src/test/java/com/google/adk/workflow/NodeRunnerTest.java @@ -37,8 +37,11 @@ import com.google.genai.types.Part; import io.reactivex.rxjava3.core.Completable; import io.reactivex.rxjava3.core.Flowable; +import io.reactivex.rxjava3.core.Single; import io.reactivex.rxjava3.observers.TestObserver; import io.reactivex.rxjava3.subjects.CompletableSubject; +import io.reactivex.rxjava3.subjects.SingleSubject; +import io.reactivex.rxjava3.subscribers.TestSubscriber; import java.util.ArrayList; import java.util.List; import org.jspecify.annotations.Nullable; @@ -54,15 +57,17 @@ public final class NodeRunnerTest { private final List sentEvents = new ArrayList<>(); private final NodeRunner.EventSink recordingSink = - event -> Completable.fromAction(() -> sentEvents.add(event)); + (event, nodeFailure) -> Completable.fromAction(() -> sentEvents.add(event)); private final List acknowledgements = new ArrayList<>(); private final NodeRunner.EventSink acknowledgingSink = - event -> { + (event, nodeFailure) -> { sentEvents.add(event); CompletableSubject acknowledgement = CompletableSubject.create(); acknowledgements.add(acknowledgement); return acknowledgement; }; + private final List eventsInProcessing = new ArrayList<>(); + private final List> processings = new ArrayList<>(); @Test public void run_stampsTheEventWithTheNodeNamePathAndInvocation() { @@ -677,6 +682,139 @@ public void build_underAParentWithAnInvocationContext_throwsIllegalStateExceptio assertThrows(IllegalStateException.class, builder::build); } + @Test + public void runRoot_asksForTheNextEventOnlyAfterTheCallerReceivesTheProcessedEvent() { + List emitted = new ArrayList<>(); + Node node = + new TestNode( + "pair", + (ctx, input) -> + Flowable.just( + newEvent().content(MESSAGE).build(), newEvent().content(MESSAGE).build()) + .doOnNext(emitted::add)); + + TestSubscriber subscriber = TestSubscriber.create(/* initialRequested= */ 0); + + NodeRunner.runRoot(node, invocationContext(/* branch= */ null), this::holdProcessing) + .subscribe(subscriber); + finishProcessing(0); + int emittedBeforeTheCallerReceivesTheFirstEvent = emitted.size(); + subscriber.request(1); + int emittedAfterTheCallerReceivesTheFirstEvent = emitted.size(); + finishProcessing(1); + subscriber.request(1); + + assertThat(emittedBeforeTheCallerReceivesTheFirstEvent).isEqualTo(1); + assertThat(emittedAfterTheCallerReceivesTheFirstEvent).isEqualTo(2); + subscriber.assertValueCount(2).assertComplete(); + } + + @Test + public void runRoot_signalsTheNodeFailureOnlyAfterItsErrorEventIsProcessed() { + IllegalStateException failure = new IllegalStateException("node failed"); + Node node = new TestNode("boom", (ctx, input) -> Flowable.error(failure)); + + TestSubscriber subscriber = + NodeRunner.runRoot(node, invocationContext(/* branch= */ null), this::holdProcessing) + .test(); + subscriber.assertEmpty(); + finishProcessing(0); + + assertThat(getOnlyElement(subscriber.values()).errorCode()).hasValue(ILLEGAL_STATE); + subscriber.assertError(failure); + } + + @Test + public void runRoot_processingTheErrorEventFails_keepsTheNodeFailureAsSuppressed() { + IllegalStateException nodeFailure = new IllegalStateException("node failed"); + IllegalArgumentException processingFailure = new IllegalArgumentException("append failed"); + Node node = new TestNode("boom", (ctx, input) -> Flowable.error(nodeFailure)); + + TestSubscriber subscriber = + NodeRunner.runRoot( + node, + invocationContext(/* branch= */ null), + event -> + event.errorCode().isPresent() + ? Single.error(processingFailure) + : Single.just(event)) + .test(); + + subscriber.assertError(processingFailure); + assertThat(processingFailure.getSuppressed()).asList().containsExactly(nodeFailure); + } + + @Test + public void runRoot_processingAPartialEventFails_keepsTheQueuedNodeFailureAsSuppressed() { + IllegalStateException nodeFailure = new IllegalStateException("node failed"); + IllegalArgumentException processingFailure = new IllegalArgumentException("plugin failed"); + Node node = + new TestNode( + "streamer", + (ctx, input) -> + Flowable.just(newEvent().content(MESSAGE).partial(true).build()) + .concatWith(Flowable.error(nodeFailure))); + + TestSubscriber subscriber = + NodeRunner.runRoot(node, invocationContext(/* branch= */ null), this::holdProcessing) + .test(); + processings.get(0).onError(processingFailure); + + subscriber.assertError(processingFailure); + assertThat(processingFailure.getSuppressed()).asList().containsExactly(nodeFailure); + } + + @Test + public void runRoot_emitsTheEventThatProcessingReturns() { + Node node = + new TestNode("solo", (ctx, input) -> Flowable.just(newEvent().content(MESSAGE).build())); + Event replacement = newEvent().author("plugin").content(MESSAGE).build(); + + TestSubscriber subscriber = + NodeRunner.runRoot( + node, invocationContext(/* branch= */ null), event -> Single.just(replacement)) + .test(); + + subscriber.assertComplete(); + assertThat(getOnlyElement(subscriber.values())).isSameInstanceAs(replacement); + } + + @Test + public void runRoot_errorFromTheNode_failsTheRunWithoutAnErrorEvent() { + AssertionError fatal = new AssertionError("fatal"); + Node node = new TestNode("fatal", (ctx, input) -> Flowable.error(fatal)); + + TestSubscriber subscriber = + NodeRunner.runRoot(node, invocationContext(/* branch= */ null), this::holdProcessing) + .test(); + + subscriber.assertError(fatal); + assertThat(eventsInProcessing).isEmpty(); + } + + @Test + public void runRoot_partialEvent_doesNotHoldBackTheNode() { + List emitted = new ArrayList<>(); + Node node = + new TestNode( + "streamer", + (ctx, input) -> + Flowable.just( + newEvent().content(MESSAGE).partial(true).build(), + newEvent().content(MESSAGE).build()) + .doOnNext(emitted::add)); + + TestSubscriber subscriber = + NodeRunner.runRoot(node, invocationContext(/* branch= */ null), this::holdProcessing) + .test(); + int emittedWhileThePartialEventIsProcessed = emitted.size(); + finishProcessing(0); + finishProcessing(1); + + assertThat(emittedWhileThePartialEventIsProcessed).isEqualTo(2); + subscriber.assertValueCount(2).assertComplete(); + } + @CanIgnoreReturnValue private Context runAtRoot(Node node) { return rootRunner(node, recordingSink).run(/* nodeInput= */ null).blockingGet(); @@ -690,6 +828,17 @@ private static NodeRunner rootRunner(Node node, NodeRunner.EventSink eventSink) .build(); } + private Single holdProcessing(Event event) { + eventsInProcessing.add(event); + SingleSubject processing = SingleSubject.create(); + processings.add(processing); + return processing; + } + + private void finishProcessing(int index) { + processings.get(index).onSuccess(eventsInProcessing.get(index)); + } + private Context runUnderParent(Node node, boolean useSubBranch, @Nullable String overrideBranch) { return NodeRunner.builder() .node(node)