From d5ff72123c96e856c52d3746ae673ca5bd687500 Mon Sep 17 00:00:00 2001 From: Damian Momot Date: Wed, 14 Oct 2026 10:52:17 -0700 Subject: [PATCH] feat: run a workflow node as the root of an invocation PiperOrigin-RevId: 999474852 --- .../java/com/google/adk/agents/BaseAgent.java | 15 +- .../java/com/google/adk/agents/Context.java | 44 +- .../google/adk/agents/InvocationContext.java | 30 + .../com/google/adk/agents/NodeViewAgent.java | 72 ++ .../main/java/com/google/adk/apps/App.java | 61 +- .../java/com/google/adk/runner/Runner.java | 164 +++- .../adk/workflow/NodeExecutionState.java | 240 ++++++ .../com/google/adk/workflow/NodeRunner.java | 470 +++++++++++ .../com/google/adk/agents/BaseAgentTest.java | 34 +- .../com/google/adk/agents/ContextTest.java | 7 + .../adk/agents/InvocationContextTest.java | 49 ++ .../google/adk/agents/NodeViewAgentTest.java | 95 +++ .../java/com/google/adk/apps/AppTest.java | 98 +++ .../google/adk/runner/RunnerRootNodeTest.java | 524 ++++++++++++ .../java/com/google/adk/testing/TestNode.java | 38 + .../adk/workflow/NodeExecutionStateTest.java | 54 ++ .../google/adk/workflow/NodeRunnerTest.java | 771 ++++++++++++++++++ 17 files changed, 2695 insertions(+), 71 deletions(-) create mode 100644 core/src/main/java/com/google/adk/agents/NodeViewAgent.java create mode 100644 core/src/main/java/com/google/adk/workflow/NodeExecutionState.java create mode 100644 core/src/main/java/com/google/adk/workflow/NodeRunner.java create mode 100644 core/src/test/java/com/google/adk/agents/NodeViewAgentTest.java create mode 100644 core/src/test/java/com/google/adk/apps/AppTest.java create mode 100644 core/src/test/java/com/google/adk/runner/RunnerRootNodeTest.java create mode 100644 core/src/test/java/com/google/adk/testing/TestNode.java create mode 100644 core/src/test/java/com/google/adk/workflow/NodeExecutionStateTest.java create mode 100644 core/src/test/java/com/google/adk/workflow/NodeRunnerTest.java diff --git a/core/src/main/java/com/google/adk/agents/BaseAgent.java b/core/src/main/java/com/google/adk/agents/BaseAgent.java index 81f735b25..25a405752 100644 --- a/core/src/main/java/com/google/adk/agents/BaseAgent.java +++ b/core/src/main/java/com/google/adk/agents/BaseAgent.java @@ -273,7 +273,8 @@ public ImmutableList canonicalAfterAgentCallbacks( */ private InvocationContext createInvocationContext(InvocationContext parentContext) { 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); // Check for branch to be truthy (not None, not empty string), parentContext .branch() @@ -489,13 +490,21 @@ private Event checkpointEvent(InvocationContext context, EventActions actions) { /** * Runs this agent as a workflow node through its {@link #runAsync} lifecycle, including its - * callbacks. The node input is not used. + * callbacks. The node input is not used. Each event's author is set with {@link + * com.google.adk.workflow.NodeExecutionState#eventAuthor(String)}, so the agent's events keep + * their own authors instead of an inherited author or the node's name. */ @Experimental @Override public Flowable runNode( com.google.adk.agents.Context context, @Nullable Object nodeInput) { - return runAsync(context.invocationContext()); + return runAsync(context.invocationContext()) + .doOnNext( + event -> { + if (!isNullOrEmpty(event.author())) { + context.requireNodeState().eventAuthor(event.author()); + } + }); } /** diff --git a/core/src/main/java/com/google/adk/agents/Context.java b/core/src/main/java/com/google/adk/agents/Context.java index 0432c97c9..a3f601d2d 100644 --- a/core/src/main/java/com/google/adk/agents/Context.java +++ b/core/src/main/java/com/google/adk/agents/Context.java @@ -16,11 +16,13 @@ package com.google.adk.agents; +import com.google.adk.annotations.Experimental; import com.google.adk.artifacts.ListArtifactsResponse; import com.google.adk.events.EventActions; import com.google.adk.events.ToolConfirmation; import com.google.adk.memory.SearchMemoryResponse; import com.google.adk.sessions.State; +import com.google.adk.workflow.NodeExecutionState; import com.google.common.base.Preconditions; import com.google.genai.types.Part; import io.reactivex.rxjava3.core.Completable; @@ -31,12 +33,13 @@ import org.jspecify.annotations.Nullable; /** - * The context passed to callbacks and tools during an invocation. + * The context passed to callbacks, tools, and workflow nodes during an invocation. * *

Agent and model callbacks receive a {@link CallbackContext}, and tools and tool callbacks a * {@link com.google.adk.tools.ToolContext}; both extend this class, so a helper that serves both, * or a function tool's {@code toolContext} parameter, can take a {@code Context}. {@link - * #requestConfirmation()} needs a function call ID, which only a tool call has. + * #requestConfirmation()} needs a function call ID, which only a tool call has, and {@link + * #requireNodeState()} needs the context of a workflow node. */ public class Context extends ReadonlyContext { @@ -45,6 +48,7 @@ public class Context extends ReadonlyContext { private final State state; private final @Nullable String eventId; + private final @Nullable NodeExecutionState nodeState; private Optional functionCallId = Optional.empty(); private Optional toolConfirmation = Optional.empty(); @@ -59,10 +63,32 @@ public class Context extends ReadonlyContext { InvocationContext invocationContext, @Nullable EventActions eventActions, @Nullable String eventId) { + this(invocationContext, eventActions, eventId, /* nodeState= */ null); + } + + /** + * Creates the context of the workflow node run whose state is {@code nodeState}. Only the + * workflow runtime creates a node state. + */ + @Experimental + public Context(InvocationContext invocationContext, NodeExecutionState nodeState) { + this( + invocationContext, + /* eventActions= */ null, + /* eventId= */ null, + Preconditions.checkNotNull(nodeState)); + } + + private Context( + InvocationContext invocationContext, + @Nullable EventActions eventActions, + @Nullable String eventId, + @Nullable NodeExecutionState nodeState) { super(invocationContext); this.eventActions = eventActions != null ? eventActions : EventActions.builder().build(); this.state = new State(invocationContext.session().state(), this.eventActions.stateDelta()); this.eventId = eventId; + this.nodeState = nodeState; } /** Returns the delta-aware state of the current context. */ @@ -223,4 +249,18 @@ public Single searchMemory(String query) { .searchMemory( invocationContext.session().appName(), invocationContext.session().userId(), query); } + + /** + * Returns the state of the workflow node run this context belongs to. + * + * @throws UnsupportedOperationException if this is not the context of a workflow node. + */ + @Experimental + public final NodeExecutionState requireNodeState() { + if (nodeState == null) { + throw new UnsupportedOperationException( + "Only the context of a workflow node has a node state."); + } + return nodeState; + } } 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..f73a78f5f --- /dev/null +++ b/core/src/main/java/com/google/adk/agents/NodeViewAgent.java @@ -0,0 +1,72 @@ +/* + * 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; + } + + @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..5b85cb232 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; @@ -220,6 +225,7 @@ public Builder uuidProvider(@Nullable UuidProvider uuidProvider) { @SuppressWarnings("deprecation") public Runner build() { BaseAgent buildAgent; + Node buildNode; String buildAppName; List buildPlugins; EventsCompactionConfig buildEventsCompactionConfig; @@ -234,6 +240,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 +248,7 @@ public Runner build() { buildResumabilityConfig = this.app.resumabilityConfig(); } else { buildAgent = this.agent; + buildNode = null; buildAppName = this.appName; buildPlugins = this.plugins; buildEventsCompactionConfig = null; @@ -251,6 +259,10 @@ public Runner build() { if (buildAgent == null) { throw new IllegalStateException("Agent must be provided via app() or agent()."); } + if (buildAgent instanceof NodeViewAgent && buildNode == null) { + throw new IllegalStateException( + "A node view agent cannot run on its own; run its node through an app's root node."); + } if (buildAppName == null) { throw new IllegalStateException("App name must be provided via app() or appName()."); } @@ -270,6 +282,7 @@ public Runner build() { } return new Runner( buildAgent, + buildNode, buildAppName, artifactService, sessionService, @@ -364,6 +377,7 @@ protected Runner( @Nullable ResumabilityConfig resumabilityConfig) { this( agent, + /* node= */ null, appName, artifactService, sessionService, @@ -379,6 +393,7 @@ protected Runner( private Runner( BaseAgent agent, + @Nullable Node node, String appName, BaseArtifactService artifactService, BaseSessionService sessionService, @@ -390,7 +405,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 +442,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 +764,10 @@ 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); + } Context capturedContext = Context.current(); BaseAgent rootAgent = this.agent; String invocationId = InvocationContext.newInvocationContextId(this.uuidProvider); @@ -844,9 +876,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 +888,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,6 +985,9 @@ 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); } @@ -959,6 +1006,12 @@ 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( + "Resuming an invocation of a workflow node root is not supported.")); + } return resumeCore(session, resolvedInvocationId, newMessage, runConfig, stateDelta); }); } @@ -1227,6 +1280,29 @@ 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."); + } + /** * 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 @@ -1380,7 +1456,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 +1514,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 +1593,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/NodeExecutionState.java b/core/src/main/java/com/google/adk/workflow/NodeExecutionState.java new file mode 100644 index 000000000..1801e52f2 --- /dev/null +++ b/core/src/main/java/com/google/adk/workflow/NodeExecutionState.java @@ -0,0 +1,240 @@ +/* + * 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.workflow; + +import static com.google.common.base.Preconditions.checkNotNull; +import static com.google.common.base.Preconditions.checkState; + +import com.google.adk.agents.Context; +import com.google.adk.annotations.Experimental; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; +import com.google.errorprone.annotations.CanIgnoreReturnValue; +import java.util.Collection; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Optional; +import java.util.Set; +import org.jspecify.annotations.Nullable; + +/** + * The state of one workflow node run, held by the run's {@link Context}. Only the workflow runtime + * creates it and updates the bookkeeping that records which results an event has carried. + */ +@Experimental +public final class NodeExecutionState { + + private final Node node; + private final NodeRunner.EventSink eventSink; + private final @Nullable Context parent; + private final String runId; + private final String nodePath; + private String eventAuthor; + private final Set interruptIds = new LinkedHashSet<>(); + private @Nullable Object output; + private boolean outputEmitted; + private @Nullable ImmutableList selectedRoutes; + private boolean routesEmitted; + private @Nullable Exception failure; + + private NodeExecutionState(Builder builder) { + checkState(builder.node != null, "A node state needs a node."); + checkState(builder.eventSink != null, "A node state needs an event sink."); + checkState(builder.runId != null, "A node state needs a run ID."); + this.node = builder.node; + this.eventSink = builder.eventSink; + this.parent = builder.parent; + this.runId = builder.runId; + NodeExecutionState parentState = parent == null ? null : parent.requireNodeState(); + this.nodePath = + NodeRunner.appendSegment( + parentState == null ? null : parentState.nodePath, + node.name(), + runId, + /* separator= */ '/'); + this.eventAuthor = parentState == null ? "" : parentState.eventAuthor; + } + + /** Returns a builder for the state of one node run. */ + static Builder builder() { + return new Builder(); + } + + /** Returns the node this run belongs to. */ + public Node node() { + return node; + } + + /** + * Returns the path of this run: {@code name@runId} segments joined by {@code /}, starting at the + * outermost node. + */ + public String nodePath() { + return nodePath; + } + + /** Returns the ID of this run among the node's runs, counting from {@code "1"}. */ + public String runId() { + return runId; + } + + /** Returns the context of the node that ran this node, or empty if this is a root node. */ + public Optional parent() { + return Optional.ofNullable(parent); + } + + /** + * Returns the author of the events this run emits, or empty to use the node's name. A run under a + * parent starts with the parent's author. + */ + public String eventAuthor() { + return eventAuthor; + } + + /** Sets the author of the events this run emits. */ + public void eventAuthor(String eventAuthor) { + this.eventAuthor = checkNotNull(eventAuthor); + } + + /** Returns the node's output, or empty if it has produced none. */ + public Optional output() { + return Optional.ofNullable(output); + } + + /** + * Records the node's output, which an event carries when the node finishes unless an earlier + * event already carried it. A node produces at most one output. + * + * @throws IllegalStateException if the node already has an output. + */ + public void produceOutput(Object output) { + checkNotNull(output); + checkState( + this.output == null, + "Node '%s' produced a second output; a node produces at most one output.", + node.name()); + this.output = output; + } + + /** + * Returns the routes the node selected, which decide the outgoing edges to follow, or empty if it + * made no routing decision. An empty list is a decision that selects no route. + */ + public Optional> selectedRoutes() { + return Optional.ofNullable(selectedRoutes); + } + + /** + * Selects the node's routes, replacing any earlier selection. An event carries them when the node + * finishes unless an earlier event already carried them. + */ + public void selectRoutes(List routes) { + this.selectedRoutes = ImmutableList.copyOf(routes); + this.routesEmitted = false; + } + + /** Returns the IDs of the long-running tool calls the node is waiting on. */ + public ImmutableSet interruptIds() { + return ImmutableSet.copyOf(interruptIds); + } + + NodeRunner.EventSink eventSink() { + return eventSink; + } + + /** Returns whether an event has carried the node's output. */ + boolean outputEmitted() { + return outputEmitted; + } + + void markOutputEmitted() { + outputEmitted = true; + } + + /** Returns whether an event has carried the node's current route selection. */ + boolean routesEmitted() { + return routesEmitted; + } + + void markRoutesEmitted() { + routesEmitted = true; + } + + void addInterruptIds(Collection ids) { + interruptIds.addAll(ids); + } + + /** Returns the exception the node failed with, if it failed. */ + Optional failure() { + return Optional.ofNullable(failure); + } + + void failure(Exception failure) { + this.failure = checkNotNull(failure); + } + + /** Builder for {@link NodeExecutionState}. */ + @Experimental + static final class Builder { + private @Nullable Node node; + private NodeRunner.@Nullable EventSink eventSink; + private @Nullable Context parent; + private @Nullable String runId; + + private Builder() {} + + /** Sets the node this run belongs to. */ + @CanIgnoreReturnValue + Builder node(Node node) { + this.node = checkNotNull(node); + return this; + } + + /** Sets the sink that receives the events of the run. */ + @CanIgnoreReturnValue + Builder eventSink(NodeRunner.EventSink eventSink) { + this.eventSink = checkNotNull(eventSink); + return this; + } + + /** + * Sets the context of the node that ran this node, or null for a root node. A run under a + * parent extends the parent's node path and starts with the parent's event author. + */ + @CanIgnoreReturnValue + Builder parent(@Nullable Context parent) { + this.parent = parent; + return this; + } + + /** Sets the ID of this run among the node's runs. */ + @CanIgnoreReturnValue + Builder runId(String runId) { + this.runId = checkNotNull(runId); + return this; + } + + /** + * Builds the state. + * + * @throws IllegalStateException if the node, the event sink, or the run ID is not set. + */ + NodeExecutionState build() { + return new NodeExecutionState(this); + } + } +} diff --git a/core/src/main/java/com/google/adk/workflow/NodeRunner.java b/core/src/main/java/com/google/adk/workflow/NodeRunner.java new file mode 100644 index 000000000..d6d69af51 --- /dev/null +++ b/core/src/main/java/com/google/adk/workflow/NodeRunner.java @@ -0,0 +1,470 @@ +/* + * 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.workflow; + +import static com.google.common.base.Preconditions.checkNotNull; +import static com.google.common.base.Preconditions.checkState; +import static com.google.common.base.Strings.isNullOrEmpty; +import static com.google.common.base.Strings.nullToEmpty; + +import com.google.adk.agents.Context; +import com.google.adk.agents.InvocationContext; +import com.google.adk.annotations.Experimental; +import com.google.adk.events.Event; +import com.google.adk.events.EventActions; +import com.google.common.collect.ImmutableList; +import com.google.errorprone.annotations.CanIgnoreReturnValue; +import com.google.genai.types.FinishReason; +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.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. {@link #runRoot} runs a node as the root of an + * invocation. + */ +@Experimental +public final class NodeRunner { + + /** Sink that receives the events of a workflow run. */ + @Experimental + 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. + */ + Completable send(Event event); + } + + private final Node node; + private final InvocationContext invocationContext; + private final EventSink eventSink; + private final @Nullable Context parent; + private final String runId; + private final boolean useSubBranch; + private final @Nullable String overrideBranch; + + private NodeRunner(Builder builder) { + checkState(builder.node != null, "A node runner needs a node."); + if (builder.parent != null) { + checkState( + builder.invocationContext == null && builder.eventSink == null, + "A node runner with a parent takes its invocation context and event sink from it."); + this.invocationContext = builder.parent.invocationContext(); + this.eventSink = builder.parent.requireNodeState().eventSink(); + } else { + checkState( + builder.invocationContext != null && builder.eventSink != null, + "A node runner without a parent needs an invocation context and an event sink."); + this.invocationContext = builder.invocationContext; + this.eventSink = builder.eventSink; + } + this.node = builder.node; + this.parent = builder.parent; + this.runId = builder.runId; + this.useSubBranch = builder.useSubBranch; + this.overrideBranch = builder.overrideBranch; + } + + /** Returns a builder for a runner of one node. */ + 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. + */ + @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(); + EventSink sink = + event -> { + CompletableSubject processed = CompletableSubject.create(); + queue.onNext(new QueuedEvent(event, processed)); + 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 + .concatMapSingle( + queued -> + processEvent + .apply(queued.event()) + .map( + processedEvent -> + new QueuedEvent(processedEvent, queued.processed()))) + // The node proceeds only after the caller receives the processed event. + .doAfterNext(queued -> queued.processed().onComplete()) + .map(QueuedEvent::event) + .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 + * Error} propagates. + */ + Single run(@Nullable Object nodeInput) { + return Single.defer( + () -> { + Context context = + new Context( + childInvocationContext(), + NodeExecutionState.builder() + .node(node) + .eventSink(eventSink) + .parent(parent) + .runId(runId) + .build()); + return new Activation(context).run(nodeInput).toSingleDefault(context); + }); + } + + private InvocationContext childInvocationContext() { + if (useSubBranch) { + String base = + overrideBranch != null ? overrideBranch : invocationContext.branch().orElse(null); + return invocationContext.toBuilder() + .branch(appendSegment(base, node.name(), runId, /* separator= */ '.')) + .build(); + } + if (overrideBranch != null) { + return invocationContext.toBuilder().branch(overrideBranch).build(); + } + return invocationContext; + } + + /** + * One run of the node, during which the node is asked for each event only after its previous + * event is processed. Events from nested nodes pass through without reapplying their routing or + * transfers. When the node finishes, a final event carries any output, routes, or deltas not + * carried by an earlier event. + */ + private final class Activation { + private final Context context; + private final NodeExecutionState nodeState; + + Activation(Context context) { + this.context = context; + this.nodeState = context.requireNodeState(); + } + + Completable run(@Nullable Object nodeInput) { + return asBaseNode(node) + .run(context, nodeInput) + // Avoid concatMapCompletable: its prefetch would let the node run ahead of the sink. + .flatMapCompletable(this::dispatch, /* delayErrors= */ true, /* maxConcurrency= */ 1) + .andThen(Completable.defer(this::flushPending)) + .onErrorResumeNext(this::onError); + } + + private Completable onError(Throwable error) { + // Only an Exception fails the node; an Error is fatal and propagates. + return error instanceof Exception exception ? fail(exception) : Completable.error(error); + } + + private Completable dispatch(Event event) { + track(event); + return send(event); + } + + /** Records the output, interrupts, and routing carried by {@code event} on the node state. */ + private void track(Event event) { + if (event.output().isPresent()) { + nodeState.produceOutput(event.output().get()); + } else if (isMessageAsOutput(event)) { + nodeState.produceOutput(event.content().get()); + } + event.longRunningToolIds().ifPresent(nodeState::addInterruptIds); + // An event of a nested node was already acted on at that node's level. + if (isNullOrEmpty(event.author()) || event.author().equals(node.name())) { + event + .actions() + .route() + .ifPresent( + route -> { + nodeState.selectRoutes(route); + nodeState.markRoutesEmitted(); + }); + event.actions().transferToAgent().ifPresent(context.eventActions()::setTransferToAgent); + } + } + + /** + * Stamps {@code event}, moves the pending deltas onto it unless it is partial, and sends it. + */ + private Completable send(Event event) { + boolean messageAsOutput = isMessageAsOutput(event); + // A message-as-output event's content is its output, so drop the repeated output field. + Event stamped = stamp(messageAsOutput ? event.toBuilder().output(null).build() : event); + Event outgoing = stamped.partial().orElse(false) ? stamped : withPendingDeltas(stamped); + return eventSink + .send(outgoing) + .doOnComplete( + () -> { + if (outgoing.output().isPresent() || messageAsOutput) { + nodeState.markOutputEmitted(); + } + }); + } + + /** Sends a final event carrying any output, routes, or deltas that no earlier event carried. */ + private Completable flushPending() { + boolean pendingOutput = nodeState.output().isPresent() && !nodeState.outputEmitted(); + boolean pendingRoute = nodeState.selectedRoutes().isPresent() && !nodeState.routesEmitted(); + EventActions actions = context.eventActions(); + boolean pendingDeltas = !actions.stateDelta().isEmpty() || !actions.artifactDelta().isEmpty(); + if (!pendingOutput && !pendingRoute && !pendingDeltas) { + return Completable.complete(); + } + Event event = + withPendingDeltas( + stamp( + newEvent() + .output(pendingOutput ? nodeState.output().get() : null) + .actions( + EventActions.builder() + .route(pendingRoute ? nodeState.selectedRoutes().get() : null) + .build()) + .build())); + return eventSink + .send(event) + .doOnComplete( + () -> { + if (pendingOutput) { + nodeState.markOutputEmitted(); + } + if (pendingRoute) { + nodeState.markRoutesEmitted(); + } + }); + } + + /** + * Sends an error event for {@code failure} with any pending deltas to preserve state changes + * made before the failure, then records the failure on the node state. + */ + private Completable fail(Exception failure) { + Event errorEvent = + withPendingDeltas( + stamp( + newEvent() + .errorCode(new FinishReason(errorTypeName(failure))) + .errorMessage(nullToEmpty(failure.getMessage())) + .build())); + return eventSink.send(errorEvent).doOnComplete(() -> nodeState.failure(failure)); + } + + /** + * Moves the context's pending deltas onto a copy of {@code event} so each is sent only once. + */ + private Event withPendingDeltas(Event event) { + Map stateDelta = context.eventActions().stateDelta(); + Map artifactDelta = context.eventActions().artifactDelta(); + if (stateDelta.isEmpty() && artifactDelta.isEmpty()) { + return event; + } + EventActions actions = event.actions().toBuilder().build(); + actions.stateDelta().putAll(stateDelta); + actions.artifactDelta().putAll(artifactDelta); + stateDelta.clear(); + artifactDelta.clear(); + return event.toBuilder().actions(actions).build(); + } + + /** Attributes {@code event} to its author and to this run of the node. */ + private Event stamp(Event event) { + InvocationContext eventContext = context.invocationContext(); + String author = nodeState.eventAuthor(); + NodeInfo nodeInfo = event.nodeInfo().orElseGet(() -> NodeInfo.builder().build()); + boolean carriesOutput = event.output().isPresent() || isMessageAsOutput(event); + return event.toBuilder() + .author(author.isEmpty() ? node.name() : author) + .invocationId(eventContext.invocationId()) + .branch(event.branch().orElse(eventContext.branch().orElse(null))) + .nodeInfo( + nodeInfo.toBuilder() + .path(nodeState.nodePath()) + .outputFor( + carriesOutput ? ImmutableList.of(nodeState.nodePath()) : nodeInfo.outputFor()) + .build()) + .build(); + } + + private Event.Builder newEvent() { + InvocationContext eventContext = context.invocationContext(); + return Event.builder() + .id(eventContext.newUuid()) + .timestamp(eventContext.now().toEpochMilli()); + } + } + + /** + * An event queued for the runner, paired with the signal that completes once the caller receives + * the event. + */ + private record QueuedEvent(Event event, CompletableSubject processed) {} + + /** 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; + return isNullOrEmpty(base) ? segment : base + separator + segment; + } + + /** Returns whether {@code event} is a final event whose content is the node's output. */ + private static boolean isMessageAsOutput(Event event) { + return !event.partial().orElse(false) + && event.content().isPresent() + && event.nodeInfo().map(NodeInfo::messageAsOutput).orElse(false); + } + + /** Returns the simple name of the failure's type, or its full name for an anonymous class. */ + private static String errorTypeName(Exception failure) { + String name = failure.getClass().getSimpleName(); + return name.isEmpty() ? failure.getClass().getName() : name; + } + + /** Adapts any node to a {@link BaseNode}, whose {@link BaseNode#run} turns values into events. */ + private static BaseNode asBaseNode(Node node) { + if (node instanceof BaseNode baseNode) { + return baseNode; + } + return new BaseNode(node.name(), node.description()) { + @Override + public Flowable runNode(Context context, @Nullable Object nodeInput) { + return node.runNode(context, nodeInput); + } + }; + } + + /** + * Builder for {@link NodeRunner}. A runner at the root of a run takes an invocation context and + * an event sink, while a runner under a parent node takes both from the parent's context. + */ + @Experimental + static final class Builder { + private @Nullable Node node; + private @Nullable InvocationContext invocationContext; + private @Nullable EventSink eventSink; + private @Nullable Context parent; + private String runId = "1"; + private boolean useSubBranch; + private @Nullable String overrideBranch; + + private Builder() {} + + /** Sets the node to run. */ + @CanIgnoreReturnValue + Builder node(Node node) { + this.node = checkNotNull(node); + return this; + } + + /** Sets the invocation context of a run at the root. */ + @CanIgnoreReturnValue + Builder invocationContext(InvocationContext invocationContext) { + this.invocationContext = checkNotNull(invocationContext); + return this; + } + + /** Sets the sink that receives the events of a run at the root. */ + @CanIgnoreReturnValue + Builder eventSink(EventSink eventSink) { + this.eventSink = checkNotNull(eventSink); + return this; + } + + /** Sets the context of the node that runs this node. */ + @CanIgnoreReturnValue + Builder parent(Context parent) { + this.parent = checkNotNull(parent); + return this; + } + + /** Sets the ID of this run among the node's runs, which defaults to {@code "1"}. */ + @CanIgnoreReturnValue + Builder runId(String runId) { + this.runId = checkNotNull(runId); + return this; + } + + /** + * Sets whether the node runs on a sub-branch of its base branch, as a node fanned out from its + * parent does. Branches isolate one parallel branch's events from a sibling's model history. + */ + @CanIgnoreReturnValue + Builder useSubBranch(boolean useSubBranch) { + this.useSubBranch = useSubBranch; + return this; + } + + /** Sets the branch that replaces the inherited one, or null to keep the inherited branch. */ + @CanIgnoreReturnValue + Builder overrideBranch(@Nullable String overrideBranch) { + this.overrideBranch = overrideBranch; + return this; + } + + /** + * Builds the runner. + * + * @throws IllegalStateException if no node is set, if neither a parent nor both an invocation + * context and an event sink are set, or if a parent is set together with either of them. + */ + NodeRunner build() { + return new NodeRunner(this); + } + } +} 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 2792fbb76..7f3f0a781 100644 --- a/core/src/test/java/com/google/adk/agents/BaseAgentTest.java +++ b/core/src/test/java/com/google/adk/agents/BaseAgentTest.java @@ -25,6 +25,7 @@ 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; @@ -408,31 +409,18 @@ public void runAsync_firstAfterCallbackReturnsContent_skipsSecondAfterCallback() } @Test - public void runNode_runsAgentWithItsCallbacks() { - var runAsyncImpl = TestCallback.returningEmpty(); - Content runAsyncImplContent = Content.fromParts(Part.fromText("main_output")); - Content afterCallbackContent = Content.fromParts(Part.fromText("after_callback_output")); - var beforeCallback = TestCallback.returningEmpty(); - var afterCallback = TestCallback.returning(afterCallbackContent); - TestBaseAgent agent = - new TestBaseAgent( - TEST_AGENT_NAME, - TEST_AGENT_DESCRIPTION, - ImmutableList.of(beforeCallback.asBeforeAgentCallback()), - ImmutableList.of(afterCallback.asAfterAgentCallback()), - runAsyncImpl.asRunAsyncImplSupplier(runAsyncImplContent)); - Context context = - new CallbackContext(TestUtils.createInvocationContext(agent), /* eventActions= */ null); - Node node = agent; + 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(); - List results = - node.runNode(context, /* nodeInput= */ "unused").cast(Event.class).toList().blockingGet(); + agent.runAsync(parentContext).blockingSubscribe(); - assertThat(results).hasSize(2); - assertThat(results.get(0).content()).hasValue(runAsyncImplContent); - assertThat(results.get(1).content()).hasValue(afterCallbackContent); - assertThat(beforeCallback.wasCalled()).isTrue(); - assertThat(runAsyncImpl.wasCalled()).isTrue(); + assertThat(agent.getLastInvocationContext().node()).isEmpty(); } @Test diff --git a/core/src/test/java/com/google/adk/agents/ContextTest.java b/core/src/test/java/com/google/adk/agents/ContextTest.java index 6ff2b4365..29f4023e4 100644 --- a/core/src/test/java/com/google/adk/agents/ContextTest.java +++ b/core/src/test/java/com/google/adk/agents/ContextTest.java @@ -76,6 +76,13 @@ public void requestConfirmation_onCallbackContext_throwsIllegalStateException() assertThrows(IllegalStateException.class, callbackContext::requestConfirmation); } + @Test + public void requireNodeState_onCallbackContext_throwsUnsupportedOperationException() { + CallbackContext context = new CallbackContext(newInvocationContext(), null); + + assertThrows(UnsupportedOperationException.class, context::requireNodeState); + } + @Test public void searchMemory_onCallbackContext_findsRememberedEvents() { InvocationContext invocationContext = newInvocationContext(); 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..9dc4cbfc2 --- /dev/null +++ b/core/src/test/java/com/google/adk/agents/NodeViewAgentTest.java @@ -0,0 +1,95 @@ +/* + * 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 runAsyncImpl_failsWithUnsupportedOperationException() { + NodeViewAgent agent = new NodeViewAgent(NODE); + + agent.runAsyncImpl(contextFor(agent)).test().assertError(UnsupportedOperationException.class); + } + + @Test + public void runLiveImpl_failsWithUnsupportedOperationException() { + NodeViewAgent agent = new NodeViewAgent(NODE); + + agent.runLiveImpl(contextFor(agent)).test().assertError(UnsupportedOperationException.class); + } + + private static InvocationContext contextFor(NodeViewAgent agent) { + return TestUtils.createInvocationContext(TestUtils.createRootAgent()).toBuilder() + .agent(agent) + .node(agent.node()) + .build(); + } +} 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..9aa930898 --- /dev/null +++ b/core/src/test/java/com/google/adk/runner/RunnerRootNodeTest.java @@ -0,0 +1,524 @@ +/* + * 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.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(); + + @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_throwsIllegalStateException() { + Runner.Builder builder = + Runner.builder().agent(new NodeViewAgent(TestNode.of("solo"))).appName(APP_NAME); + + assertThrows(IllegalStateException.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 = + runner( + App.builder() + .name(APP_NAME) + .rootNode(node) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build()); + 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); + + IllegalArgumentException error = + assertThrows( + IllegalArgumentException.class, + () -> runner.runAsync(USER_ID, session.id(), message).toList().blockingGet()); + + 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 = + runner( + App.builder() + .name(APP_NAME) + .rootNode(countingNode(runs)) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build()); + Session session = newSession(runner); + Content message = textWithFunctionResponse(IDENTIFIED_RESPONSE); + + IllegalArgumentException error = + assertThrows( + IllegalArgumentException.class, + () -> runner.runAsync(USER_ID, session.id(), message).toList().blockingGet()); + + 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); + Content message = + textWithFunctionResponse( + FunctionResponse.builder().name("tool").response(ImmutableMap.of("v", 1)).build()); + + runner.runAsync(USER_ID, session.id(), message).blockingSubscribe(); + + 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 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 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 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/testing/TestNode.java b/core/src/test/java/com/google/adk/testing/TestNode.java new file mode 100644 index 000000000..9bf8c9979 --- /dev/null +++ b/core/src/test/java/com/google/adk/testing/TestNode.java @@ -0,0 +1,38 @@ +/* + * 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 com.google.adk.agents.Context; +import com.google.adk.workflow.Node; +import io.reactivex.rxjava3.core.Flowable; +import java.util.function.BiFunction; +import org.jspecify.annotations.Nullable; + +/** A workflow node whose run is given by a function, so a test can declare its node inline. */ +public record TestNode(String name, BiFunction> body) + implements Node { + + /** Returns a node named {@code name} that emits nothing. */ + public static TestNode of(String name) { + return new TestNode(name, (context, nodeInput) -> Flowable.empty()); + } + + @Override + public Flowable runNode(Context context, @Nullable Object nodeInput) { + return body.apply(context, nodeInput); + } +} diff --git a/core/src/test/java/com/google/adk/workflow/NodeExecutionStateTest.java b/core/src/test/java/com/google/adk/workflow/NodeExecutionStateTest.java new file mode 100644 index 000000000..9ad3e60c9 --- /dev/null +++ b/core/src/test/java/com/google/adk/workflow/NodeExecutionStateTest.java @@ -0,0 +1,54 @@ +/* + * 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.workflow; + +import static org.junit.Assert.assertThrows; + +import com.google.adk.testing.TestNode; +import io.reactivex.rxjava3.core.Completable; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public final class NodeExecutionStateTest { + + private static final NodeRunner.EventSink SINK = event -> Completable.complete(); + + @Test + public void build_withoutANode_throwsIllegalStateException() { + NodeExecutionState.Builder builder = NodeExecutionState.builder().eventSink(SINK).runId("1"); + + assertThrows(IllegalStateException.class, builder::build); + } + + @Test + public void build_withoutAnEventSink_throwsIllegalStateException() { + NodeExecutionState.Builder builder = + NodeExecutionState.builder().node(TestNode.of("worker")).runId("1"); + + assertThrows(IllegalStateException.class, builder::build); + } + + @Test + public void build_withoutARunId_throwsIllegalStateException() { + NodeExecutionState.Builder builder = + NodeExecutionState.builder().node(TestNode.of("worker")).eventSink(SINK); + + assertThrows(IllegalStateException.class, builder::build); + } +} diff --git a/core/src/test/java/com/google/adk/workflow/NodeRunnerTest.java b/core/src/test/java/com/google/adk/workflow/NodeRunnerTest.java new file mode 100644 index 000000000..43a6bdcf2 --- /dev/null +++ b/core/src/test/java/com/google/adk/workflow/NodeRunnerTest.java @@ -0,0 +1,771 @@ +/* + * 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.workflow; + +import static com.google.common.collect.ImmutableList.toImmutableList; +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.agents.Context; +import com.google.adk.agents.InvocationContext; +import com.google.adk.events.Event; +import com.google.adk.events.EventActions; +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.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; +import com.google.errorprone.annotations.CanIgnoreReturnValue; +import com.google.genai.types.Content; +import com.google.genai.types.FinishReason; +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 java.util.function.Function; +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 NodeRunnerTest { + + private static final Content MESSAGE = Content.fromParts(Part.fromText("hi")); + private static final FinishReason ILLEGAL_STATE = new FinishReason("IllegalStateException"); + + private final List sentEvents = new ArrayList<>(); + private final NodeRunner.EventSink recordingSink = + event -> Completable.fromAction(() -> sentEvents.add(event)); + private final List acknowledgements = new ArrayList<>(); + private final NodeRunner.EventSink acknowledgingSink = + event -> { + sentEvents.add(event); + CompletableSubject acknowledgement = CompletableSubject.create(); + acknowledgements.add(acknowledgement); + return acknowledgement; + }; + private final List eventsInProcessing = new ArrayList<>(); + private final List> processings = new ArrayList<>(); + private final Function> heldProcessing = + event -> { + eventsInProcessing.add(event); + SingleSubject processing = SingleSubject.create(); + processings.add(processing); + return processing; + }; + + @Test + public void run_stampsTheEventWithTheNodeNamePathAndInvocation() { + runAtRoot(new TestNode("greet", (ctx, input) -> Flowable.just("hi"))); + + Event event = getOnlyElement(sentEvents); + assertThat(event.author()).isEqualTo("greet"); + assertThat(event.invocationId()).isEqualTo("invocationId"); + assertThat(event.nodeInfo().map(NodeInfo::path)).hasValue("greet@1"); + } + + @Test + public void run_yieldedValue_becomesTheEventAndNodeStateOutput() { + Context context = runAtRoot(new TestNode("greet", (ctx, input) -> Flowable.just("hi"))); + + Event event = getOnlyElement(sentEvents); + assertThat(event.output()).hasValue("hi"); + assertThat(event.nodeInfo().map(NodeInfo::outputFor)).hasValue(ImmutableList.of("greet@1")); + assertThat(context.requireNodeState().output()).hasValue("hi"); + } + + @Test + public void run_passesTheInputToTheNode() { + List inputs = new ArrayList<>(); + Node node = + new TestNode( + "echo", + (ctx, input) -> { + inputs.add(input); + return Flowable.empty(); + }); + + rootRunner(node, recordingSink).run("question").blockingSubscribe(); + + assertThat(inputs).containsExactly("question"); + } + + @Test + public void run_withoutABranchDirective_keepsTheParentBranch() { + Context context = + runUnderParent( + new TestNode("greet", (ctx, input) -> Flowable.just("hi")), + /* useSubBranch= */ false, + /* overrideBranch= */ null); + + assertThat(context.branch()).hasValue("root"); + assertThat(getOnlyElement(sentEvents).branch()).hasValue("root"); + } + + @Test + public void run_onASubBranch_derivesItFromTheParentBranch() { + Context context = + runUnderParent( + new TestNode("greet", (ctx, input) -> Flowable.just("hi")), + /* useSubBranch= */ true, + /* overrideBranch= */ null); + + assertThat(context.branch()).hasValue("root.greet@1"); + assertThat(getOnlyElement(sentEvents).branch()).hasValue("root.greet@1"); + } + + @Test + public void run_withAnOverrideBranch_replacesTheParentBranch() { + Context context = + runUnderParent( + new TestNode("greet", (ctx, input) -> Flowable.just("hi")), + /* useSubBranch= */ false, + /* overrideBranch= */ "merged"); + + assertThat(context.branch()).hasValue("merged"); + assertThat(getOnlyElement(sentEvents).branch()).hasValue("merged"); + } + + @Test + public void run_underAParent_extendsItsPathWithTheRunId() { + Context context = runWorkerUnder(parentContext()); + + assertThat(getOnlyElement(sentEvents).nodeInfo().map(NodeInfo::path)) + .hasValue("root@1/worker@2"); + assertThat(context.requireNodeState().runId()).isEqualTo("2"); + } + + @Test + public void run_underAParent_linksTheContextToTheParent() { + Context parent = parentContext(); + + Context context = runWorkerUnder(parent); + + assertThat(context.requireNodeState().parent()).hasValue(parent); + } + + @Test + public void run_underAParent_inheritsItsEventAuthor() { + Context parent = parentContext(); + parent.requireNodeState().eventAuthor("custom-agent"); + + runWorkerUnder(parent); + + assertThat(getOnlyElement(sentEvents).author()).isEqualTo("custom-agent"); + } + + @Test + public void run_messageAsOutputEvent_makesItsContentTheOutputWithoutAnotherEvent() { + Context context = + runAtRoot( + new TestNode( + "greet", (ctx, input) -> Flowable.just(messageAsOutput(/* output= */ null)))); + + Event event = getOnlyElement(sentEvents); + assertThat(context.requireNodeState().output()).hasValue(MESSAGE); + assertThat(event.content()).hasValue(MESSAGE); + assertThat(event.output()).isEmpty(); + assertThat(event.nodeInfo().map(NodeInfo::outputFor)).hasValue(ImmutableList.of("greet@1")); + } + + @Test + public void run_messageAsOutputEventWithAnOutputField_dropsTheField() { + runAtRoot(new TestNode("greet", (ctx, input) -> Flowable.just(messageAsOutput("hi")))); + + Event event = getOnlyElement(sentEvents); + assertThat(event.content()).hasValue(MESSAGE); + assertThat(event.output()).isEmpty(); + } + + @Test + public void run_routeOnTheNodesEvent_isRecordedWithoutAnotherEvent() { + ImmutableList route = ImmutableList.of(new Route.Tag("next")); + + Context context = + runAtRoot( + new TestNode( + "router", + (ctx, input) -> + Flowable.just( + newEvent().actions(EventActions.builder().route(route).build()).build()))); + + assertThat(context.requireNodeState().selectedRoutes()).hasValue(route); + assertThat(getOnlyElement(sentEvents).actions().route()).hasValue(route); + } + + @Test + public void run_routeSelectedOnTheNodeState_isSentOnALastEvent() { + ImmutableList route = ImmutableList.of(new Route.Tag("next")); + + runAtRoot( + new TestNode( + "router", + (ctx, input) -> { + ctx.requireNodeState().selectRoutes(route); + return Flowable.empty(); + })); + + assertThat(getOnlyElement(sentEvents).actions().route()).hasValue(route); + } + + @Test + public void run_outputProducedOnTheNodeState_isSentOnALastEvent() { + runAtRoot( + new TestNode( + "worker", + (ctx, input) -> { + ctx.requireNodeState().produceOutput("result"); + return Flowable.empty(); + })); + + Event event = getOnlyElement(sentEvents); + assertThat(event.output()).hasValue("result"); + assertThat(event.nodeInfo().map(NodeInfo::outputFor)).hasValue(ImmutableList.of("worker@1")); + } + + @Test + public void run_stateChangeWithoutAnEvent_isSentOnALastEvent() { + runAtRoot( + new TestNode( + "writer", + (ctx, input) -> { + ctx.state().put("key", "value"); + return Flowable.empty(); + })); + + assertThat(getOnlyElement(sentEvents).actions().stateDelta()).containsExactly("key", "value"); + } + + @Test + public void run_stateChange_ridesTheNextFinalEventOnly() { + runAtRoot( + new TestNode( + "writer", + (ctx, input) -> + Flowable.defer( + () -> { + ctx.state().put("key", "value"); + return Flowable.just( + newEvent().content(MESSAGE).partial(true).build(), + newEvent().content(MESSAGE).build()); + }))); + + assertThat(sentEvents).hasSize(2); + assertThat(sentEvents.get(0).actions().stateDelta()).isEmpty(); + assertThat(sentEvents.get(1).actions().stateDelta()).containsExactly("key", "value"); + } + + @Test + public void run_failure_sendsAnErrorEventWithThePendingStateChangeOnly() { + IllegalStateException failure = new IllegalStateException("node failed"); + + Context context = + runAtRoot( + new TestNode( + "boom", + (ctx, input) -> { + ctx.requireNodeState().produceOutput("unsent"); + ctx.requireNodeState().selectRoutes(ImmutableList.of(new Route.Tag("unsent"))); + ctx.state().put("key", "value"); + return Flowable.error(failure); + })); + + Event errorEvent = getOnlyElement(sentEvents); + assertThat(errorEvent.author()).isEqualTo("boom"); + assertThat(errorEvent.nodeInfo().map(NodeInfo::path)).hasValue("boom@1"); + assertThat(errorEvent.errorCode()).hasValue(ILLEGAL_STATE); + assertThat(errorEvent.errorMessage()).hasValue("node failed"); + assertThat(errorEvent.actions().stateDelta()).containsExactly("key", "value"); + assertThat(errorEvent.output()).isEmpty(); + assertThat(errorEvent.actions().route()).isEmpty(); + assertThat(context.requireNodeState().failure()).hasValue(failure); + } + + @Test + public void run_artifactChangeWithoutAnEvent_isSentOnALastEvent() { + runAtRoot( + new TestNode( + "writer", + (ctx, input) -> { + ctx.eventActions().artifactDelta().put("report.txt", 1); + return Flowable.empty(); + })); + + assertThat(getOnlyElement(sentEvents).actions().artifactDelta()) + .containsExactly("report.txt", 1); + } + + @Test + public void run_artifactChange_ridesTheNextFinalEventOnly() { + runAtRoot( + new TestNode( + "writer", + (ctx, input) -> + Flowable.defer( + () -> { + ctx.eventActions().artifactDelta().put("report.txt", 1); + return Flowable.just( + newEvent().content(MESSAGE).partial(true).build(), + newEvent().content(MESSAGE).build()); + }))); + + assertThat(sentEvents).hasSize(2); + assertThat(sentEvents.get(0).actions().artifactDelta()).isEmpty(); + assertThat(sentEvents.get(1).actions().artifactDelta()).containsExactly("report.txt", 1); + } + + @Test + public void run_failure_sendsAnErrorEventWithThePendingArtifactChange() { + runAtRoot( + new TestNode( + "boom", + (ctx, input) -> { + ctx.eventActions().artifactDelta().put("report.txt", 1); + return Flowable.error(new IllegalStateException("node failed")); + })); + + Event errorEvent = getOnlyElement(sentEvents); + assertThat(errorEvent.errorCode()).hasValue(ILLEGAL_STATE); + assertThat(errorEvent.actions().artifactDelta()).containsExactly("report.txt", 1); + } + + @Test + public void run_failureAfterEvents_sendsTheErrorEventOnceTheEventsAreProcessed() { + Node node = + new TestNode( + "boom", + (ctx, input) -> + Flowable.just( + newEvent().content(MESSAGE).build(), newEvent().content(MESSAGE).build()) + .concatWith(Flowable.error(new IllegalStateException("node failed")))); + + TestObserver observer = + rootRunner(node, acknowledgingSink).run(/* nodeInput= */ null).test(); + int sentBeforeAcknowledgements = sentEvents.size(); + acknowledgements.get(0).onComplete(); + acknowledgements.get(1).onComplete(); + acknowledgements.get(2).onComplete(); + + assertThat(sentBeforeAcknowledgements).isEqualTo(1); + assertThat(sentEvents).hasSize(3); + assertThat(sentEvents.get(2).errorCode()).hasValue(ILLEGAL_STATE); + observer.assertComplete(); + } + + @Test + public void run_secondOutput_failsTheNode() { + Context context = runAtRoot(new TestNode("twice", (ctx, input) -> Flowable.just("one", "two"))); + + assertThat(sentEvents).hasSize(2); + assertThat(sentEvents.get(0).output()).hasValue("one"); + assertThat(sentEvents.get(1).errorCode()).hasValue(ILLEGAL_STATE); + assertThat(context.requireNodeState().failure()).isPresent(); + } + + @Test + public void run_error_propagatesWithoutAnErrorEvent() { + rootRunner( + new TestNode("fatal", (ctx, input) -> Flowable.error(new AssertionError("fatal"))), + recordingSink) + .run(/* nodeInput= */ null) + .test() + .assertError(AssertionError.class); + + assertThat(sentEvents).isEmpty(); + } + + @Test + public void run_transferOnTheNodesEvent_isRecordedOnItsActions() { + Context context = + runAtRoot( + new TestNode( + "router", + (ctx, input) -> + Flowable.just( + newEvent() + .actions(EventActions.builder().transferToAgent("specialist").build()) + .build()))); + + assertThat(context.eventActions().transferToAgent()).hasValue("specialist"); + } + + @Test + public void run_eventOfANestedNode_neitherRoutesNorTransfersTheNode() { + Context context = + runAtRoot( + new TestNode( + "outer", + (ctx, input) -> + Flowable.just( + newEvent() + .author("inner") + .actions( + EventActions.builder() + .route(ImmutableList.of(new Route.Tag("next"))) + .transferToAgent("specialist") + .build()) + .build()))); + + assertThat(context.requireNodeState().selectedRoutes()).isEmpty(); + assertThat(context.eventActions().transferToAgent()).isEmpty(); + } + + @Test + public void run_routeSelectedAfterAnEventCarriedOne_isSentOnALastEvent() { + ImmutableList second = ImmutableList.of(new Route.Tag("second")); + runAtRoot( + new TestNode( + "router", + (ctx, input) -> + Flowable.just( + newEvent() + .actions( + EventActions.builder() + .route(ImmutableList.of(new Route.Tag("first"))) + .build()) + .build()) + .concatWith( + Flowable.defer( + () -> { + ctx.requireNodeState().selectRoutes(second); + return Flowable.empty(); + })))); + + assertThat(sentEvents).hasSize(2); + assertThat(sentEvents.get(1).actions().route()).hasValue(second); + } + + @Test + public void run_eventWithItsOwnBranch_keepsItsBranch() { + Context context = + runUnderParent( + new TestNode( + "worker", + (ctx, input) -> Flowable.just(newEvent().branch("own").content(MESSAGE).build())), + /* useSubBranch= */ false, + /* overrideBranch= */ null); + + assertThat(context.invocationContext().branch()).hasValue("root"); + assertThat(getOnlyElement(sentEvents).branch()).hasValue("own"); + } + + @Test + public void run_longRunningToolCalls_becomeInterruptIds() { + Context context = + runAtRoot( + new TestNode( + "asker", + (ctx, input) -> + Flowable.just( + newEvent().longRunningToolIds(ImmutableSet.of("call-1")).build()))); + + assertThat(context.requireNodeState().interruptIds()).containsExactly("call-1"); + } + + @Test + public void run_sendsEachEventAfterThePreviousOneIsProcessed() { + Node node = + new TestNode( + "pair", + (ctx, input) -> + Flowable.just( + newEvent().content(MESSAGE).build(), newEvent().content(MESSAGE).build())); + + TestObserver observer = + rootRunner(node, acknowledgingSink).run(/* nodeInput= */ null).test(); + int sentBeforeFirstAcknowledgement = sentEvents.size(); + acknowledgements.get(0).onComplete(); + int sentBeforeSecondAcknowledgement = sentEvents.size(); + observer.assertNotComplete(); + acknowledgements.get(1).onComplete(); + + assertThat(sentBeforeFirstAcknowledgement).isEqualTo(1); + assertThat(sentBeforeSecondAcknowledgement).isEqualTo(2); + observer.assertComplete(); + } + + @Test + public void run_agentNode_keepsTheAuthorOfTheAgentsEvents() { + // Agent name differs from emitted author, so stamping with the node name would fail. + Node agent = + TestUtils.createSubAgent( + "router", Event.builder().id("e1").author("worker").content(MESSAGE).build()); + + runAtRoot(agent); + + Event event = getOnlyElement(sentEvents); + assertThat(event.author()).isEqualTo("worker"); + assertThat(event.nodeInfo().map(NodeInfo::path)).hasValue("router@1"); + } + + @Test + public void run_agentNodeWithSeveralAuthors_stampsEachEventWithItsOwnAuthor() { + Node agent = + TestUtils.createSubAgent( + "team", + Event.builder().id("e1").author("planner").content(MESSAGE).build(), + Event.builder().id("e2").author("worker").content(MESSAGE).build(), + Event.builder().id("e3").author("critic").content(MESSAGE).build()); + + TestObserver observer = + rootRunner(agent, acknowledgingSink).run(/* nodeInput= */ null).test(); + acknowledgements.get(0).onComplete(); + acknowledgements.get(1).onComplete(); + acknowledgements.get(2).onComplete(); + + observer.assertComplete(); + assertThat(sentEvents.stream().map(Event::author).collect(toImmutableList())) + .containsExactly("planner", "worker", "critic") + .inOrder(); + } + + @Test + public void run_agentNode_runsTheAgentWithItsCallbacks() { + var runAsyncImpl = TestCallback.returningEmpty(); + Content afterCallbackContent = Content.fromParts(Part.fromText("after_callback_output")); + var beforeCallback = TestCallback.returningEmpty(); + var afterCallback = TestCallback.returning(afterCallbackContent); + TestBaseAgent agent = + new TestBaseAgent( + "agent", + "An agent.", + ImmutableList.of(beforeCallback.asBeforeAgentCallback()), + ImmutableList.of(afterCallback.asAfterAgentCallback()), + runAsyncImpl.asRunAsyncImplSupplier(MESSAGE)); + + runAtRoot(agent); + + assertThat(sentEvents).hasSize(2); + assertThat(sentEvents.get(0).content()).hasValue(MESSAGE); + assertThat(sentEvents.get(1).content()).hasValue(afterCallbackContent); + assertThat(beforeCallback.wasCalled()).isTrue(); + assertThat(runAsyncImpl.wasCalled()).isTrue(); + } + + @Test + public void run_agentNodeEventWithoutAnAuthor_keepsThePreviousAuthor() { + Node agent = + TestUtils.createSubAgent( + "router", + Event.builder().id("e1").author("worker").content(MESSAGE).build(), + Event.builder().id("e2").author("").content(MESSAGE).build()); + + runAtRoot(agent); + + assertThat(sentEvents.stream().map(Event::author).collect(toImmutableList())) + .containsExactly("worker", "worker") + .inOrder(); + } + + @Test + public void build_withoutANode_throwsIllegalStateException() { + NodeRunner.Builder builder = + NodeRunner.builder() + .invocationContext(invocationContext(/* branch= */ null)) + .eventSink(recordingSink); + + assertThrows(IllegalStateException.class, builder::build); + } + + @Test + public void build_atTheRootWithoutAnEventSink_throwsIllegalStateException() { + NodeRunner.Builder builder = + NodeRunner.builder() + .node(TestNode.of("worker")) + .invocationContext(invocationContext(/* branch= */ null)); + + assertThrows(IllegalStateException.class, builder::build); + } + + @Test + public void build_underAParentWithAnInvocationContext_throwsIllegalStateException() { + NodeRunner.Builder builder = + NodeRunner.builder() + .node(TestNode.of("worker")) + .parent(parentContext()) + .invocationContext(invocationContext(/* branch= */ null)); + + 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), heldProcessing) + .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), heldProcessing).test(); + subscriber.assertEmpty(); + finishProcessing(0); + + assertThat(getOnlyElement(subscriber.values()).errorCode()).hasValue(ILLEGAL_STATE); + subscriber.assertError(failure); + } + + @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), heldProcessing).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), heldProcessing).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(); + } + + private static NodeRunner rootRunner(Node node, NodeRunner.EventSink eventSink) { + return NodeRunner.builder() + .node(node) + .invocationContext(invocationContext(/* branch= */ null)) + .eventSink(eventSink) + .build(); + } + + 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) + .parent(parentContext()) + .useSubBranch(useSubBranch) + .overrideBranch(overrideBranch) + .build() + .run(/* nodeInput= */ null) + .blockingGet(); + } + + @CanIgnoreReturnValue + private Context runWorkerUnder(Context parent) { + return NodeRunner.builder() + .node(new TestNode("worker", (ctx, input) -> Flowable.just("done"))) + .parent(parent) + .runId("2") + .build() + .run(/* nodeInput= */ null) + .blockingGet(); + } + + /** Returns the context of a node named {@code root} that runs on the {@code root} branch. */ + private Context parentContext() { + return new Context( + invocationContext("root"), + NodeExecutionState.builder() + .node(TestNode.of("root")) + .eventSink(recordingSink) + .runId("1") + .build()); + } + + private static InvocationContext invocationContext(@Nullable String branch) { + return TestUtils.createInvocationContext(TestUtils.createRootAgent()).toBuilder() + .branch(branch) + .build(); + } + + private static Event.Builder newEvent() { + return Event.builder().id(Event.generateEventId()).author(""); + } + + private static Event messageAsOutput(@Nullable Object output) { + return newEvent() + .content(MESSAGE) + .output(output) + .nodeInfo(NodeInfo.builder().messageAsOutput(true).build()) + .build(); + } +}