From 73a096881a1e854abfdd7e8a0502e92538f0cea9 Mon Sep 17 00:00:00 2001 From: Damian Momot Date: Thu, 15 Oct 2026 11:58:24 -0700 Subject: [PATCH] feat: run a workflow graph PiperOrigin-RevId: 1000294791 --- .../java/com/google/adk/agents/BaseAgent.java | 19 +- .../java/com/google/adk/agents/Context.java | 42 +- .../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 | 215 +++- .../adk/workflow/NodeExecutionState.java | 269 +++++ .../com/google/adk/workflow/NodeRunner.java | 501 ++++++++++ .../com/google/adk/workflow/NodeState.java | 35 + .../com/google/adk/workflow/NodeStatus.java | 32 + .../com/google/adk/workflow/Scheduler.java | 334 +++++++ .../java/com/google/adk/workflow/Trigger.java | 35 + .../com/google/adk/workflow/TriggerQueue.java | 58 ++ .../com/google/adk/workflow/Workflow.java | 18 +- .../WorkflowConfigurationException.java} | 15 +- .../com/google/adk/agents/BaseAgentTest.java | 54 +- .../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 | 686 +++++++++++++ .../java/com/google/adk/testing/TestNode.java | 38 + .../com/google/adk/workflow/EdgeTest.java | 5 +- .../google/adk/workflow/EdgesBuilderTest.java | 17 +- .../com/google/adk/workflow/GraphTest.java | 91 +- .../adk/workflow/NodeExecutionStateTest.java | 54 + .../google/adk/workflow/NodeRunnerTest.java | 923 ++++++++++++++++++ .../google/adk/workflow/SchedulerTest.java | 651 ++++++++++++ .../com/google/adk/workflow/WorkflowTest.java | 25 +- 29 files changed, 4372 insertions(+), 157 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/main/java/com/google/adk/workflow/NodeState.java create mode 100644 core/src/main/java/com/google/adk/workflow/NodeStatus.java create mode 100644 core/src/main/java/com/google/adk/workflow/Scheduler.java create mode 100644 core/src/main/java/com/google/adk/workflow/Trigger.java create mode 100644 core/src/main/java/com/google/adk/workflow/TriggerQueue.java rename core/src/{test/java/com/google/adk/workflow/StubNode.java => main/java/com/google/adk/workflow/WorkflowConfigurationException.java} (63%) 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 create mode 100644 core/src/test/java/com/google/adk/workflow/SchedulerTest.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..21de2e577 100644 --- a/core/src/main/java/com/google/adk/agents/BaseAgent.java +++ b/core/src/main/java/com/google/adk/agents/BaseAgent.java @@ -33,6 +33,7 @@ import com.google.adk.telemetry.Tracing; import com.google.adk.utils.AgentEnums.AgentOrigin; import com.google.adk.workflow.BaseNode; +import com.google.adk.workflow.NodeExecutionState; import com.google.common.collect.ImmutableList; import com.google.errorprone.annotations.CanIgnoreReturnValue; import com.google.errorprone.annotations.DoNotCall; @@ -273,7 +274,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 +491,24 @@ 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, and each event's author is copied into the node state so + * that the agent's events keep their own authors instead of an inherited author or the node's + * name. + * + * @throws IllegalStateException if {@code context} is not the context of a workflow node. */ @Experimental @Override public Flowable runNode( com.google.adk.agents.Context context, @Nullable Object nodeInput) { - return runAsync(context.invocationContext()); + NodeExecutionState nodeState = context.requireNodeState(); + return runAsync(context.invocationContext()) + .doOnNext( + event -> { + if (!isNullOrEmpty(event.author())) { + nodeState.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..945bd5a5b 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,16 @@ 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 IllegalStateException if this is not the context of a workflow node. + */ + @Experimental + public final NodeExecutionState requireNodeState() { + Preconditions.checkState( + nodeState != null, "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..6f4f1c616 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,10 @@ import com.google.adk.tools.BaseTool; import com.google.adk.tools.FunctionTool; import com.google.adk.utils.CollectionUtils; +import com.google.adk.workflow.Edge; +import com.google.adk.workflow.Node; +import com.google.adk.workflow.NodeRunner; +import com.google.adk.workflow.Workflow; import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableList; import com.google.common.collect.MapMaker; @@ -78,6 +83,7 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.HashSet; +import java.util.LinkedHashSet; import java.util.List; import java.util.Map; import java.util.Objects; @@ -85,11 +91,13 @@ import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; +import java.util.function.Function; import org.jspecify.annotations.Nullable; /** The main class for the GenAI Agents runner. */ public class Runner { private final BaseAgent agent; + private final @Nullable Node node; private final String appName; private final BaseArtifactService artifactService; private final BaseSessionService sessionService; @@ -103,6 +111,8 @@ public class Runner { private final UuidProvider uuidProvider; private final ConcurrentMap activeSessionCompletables = new MapMaker().weakValues().makeMap(); + private static final String NODE_ROOT_RESUME_UNSUPPORTED = + "Resuming an invocation of a workflow node root is not supported."; /** Builder for {@link Runner}. */ public static class Builder { @@ -220,6 +230,7 @@ public Builder uuidProvider(@Nullable UuidProvider uuidProvider) { @SuppressWarnings("deprecation") public Runner build() { BaseAgent buildAgent; + Node buildNode; String buildAppName; List buildPlugins; EventsCompactionConfig buildEventsCompactionConfig; @@ -234,6 +245,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 +253,7 @@ public Runner build() { buildResumabilityConfig = this.app.resumabilityConfig(); } else { buildAgent = this.agent; + buildNode = null; buildAppName = this.appName; buildPlugins = this.plugins; buildEventsCompactionConfig = null; @@ -270,6 +283,7 @@ public Runner build() { } return new Runner( buildAgent, + buildNode, buildAppName, artifactService, sessionService, @@ -364,6 +378,7 @@ protected Runner( @Nullable ResumabilityConfig resumabilityConfig) { this( agent, + /* node= */ null, appName, artifactService, sessionService, @@ -379,6 +394,7 @@ protected Runner( private Runner( BaseAgent agent, + @Nullable Node node, String appName, BaseArtifactService artifactService, BaseSessionService sessionService, @@ -390,7 +406,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 +443,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; } @@ -448,10 +477,27 @@ public PluginManager pluginManager() { public Completable close() { List completables = new ArrayList<>(); completables.add(agent.close()); + if (node != null) { + Set nodeAgents = new LinkedHashSet<>(); + collectAgents(node, nodeAgents); + nodeAgents.forEach(nodeAgent -> completables.add(nodeAgent.close())); + } completables.add(this.pluginManager.close()); return Completable.mergeDelayError(completables); } + /** Adds the agents that {@code node} runs, through its own graph and nested workflows. */ + private static void collectAgents(Node node, Set agents) { + if (node instanceof BaseAgent agent) { + agents.add(agent); + } else if (node instanceof Workflow workflow) { + for (Edge edge : workflow.edges()) { + collectAgents(edge.from(), agents); + collectAgents(edge.to(), agents); + } + } + } + /** * Prepares and creates the user-authored {@link Event}, saving any input blobs as artifacts and * attaching state deltas if configured. @@ -736,6 +782,11 @@ private Flowable runNewInvocation( @Nullable Map stateDelta) { return Flowable.defer( () -> { + // Rejected before any callback or append; defer turns the throw into an error signal. + if (node != null) { + checkNoTextWithFunctionResponses(newMessage); + checkNoAnswerToAnEarlierInvocation(session, newMessage); + } Context capturedContext = Context.current(); BaseAgent rootAgent = this.agent; String invocationId = InvocationContext.newInvocationContextId(this.uuidProvider); @@ -844,9 +895,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 +907,59 @@ private Flowable executeAgentPipeline( // Let BaseLlmFlow block each step until this Runner has persisted the prior step's events. PersistBarrier.enable(context); - Flowable agentEvents = - context - .agent() - .runAsync(context) - .concatMap( - agentEvent -> { - // Mirror ADK Python (runners.py): partial events are streamed to the caller but - // never persisted, so managed session services (e.g. VertexAiSessionService) do - // not store a duplicate of the function call/text that the final aggregated event - // already carries. Nothing to persist, so resolve the barrier immediately. - Single persistStep = - agentEvent.partial().orElse(false) - ? Single.just(agentEvent) - : this.sessionService.appendEvent(sessionToPersist, agentEvent); - return persistStep - // Release (or fail) BaseLlmFlow's wait for this step; the Runner stays the - // sole appendEvent caller (see PersistBarrier). - .doOnSuccess( - unusedEvent -> PersistBarrier.markPersisted(context, agentEvent.id())) - .doOnError( - error -> PersistBarrier.markFailed(context, agentEvent.id(), error)) - .flatMap( - registeredEvent -> { - if (onEachPersisted != null) { - onEachPersisted.run(); - } - return context - .pluginManager() - .onEventCallback(context, registeredEvent) - .defaultIfEmpty(registeredEvent); - }) - .toFlowable(); - }); - return beforeRunEvent .toFlowable() - .switchIfEmpty(agentEvents) + .switchIfEmpty( + runRoot( + context, event -> processEvent(context, sessionToPersist, onEachPersisted, event))) .concatWith(Completable.defer(() -> pluginManager.afterRunCallback(context))) .concatWith(Completable.defer(() -> compactEvents(sessionToPersist))); } + /** + * Persists {@code event} unless it is partial, runs {@code onEachPersisted} and the plugins' + * event callbacks, and returns the event to emit. + */ + private Single processEvent( + InvocationContext context, + Session sessionToPersist, + @Nullable Runnable onEachPersisted, + Event event) { + // Mirror ADK Python (runners.py): partial events are streamed to the caller but + // never persisted, so managed session services (e.g. VertexAiSessionService) do + // not store a duplicate of the function call/text that the final aggregated event + // already carries. Nothing to persist, so resolve the barrier immediately. + Single persistStep = + event.partial().orElse(false) + ? Single.just(event) + : this.sessionService.appendEvent(sessionToPersist, event); + return persistStep + // Release (or fail) BaseLlmFlow's wait for this step; the Runner stays the + // sole appendEvent caller (see PersistBarrier). + .doOnSuccess(unusedEvent -> PersistBarrier.markPersisted(context, event.id())) + .doOnError(error -> PersistBarrier.markFailed(context, event.id(), error)) + .flatMap( + registeredEvent -> { + if (onEachPersisted != null) { + onEachPersisted.run(); + } + return context + .pluginManager() + .onEventCallback(context, registeredEvent) + .defaultIfEmpty(registeredEvent); + }); + } + + /** Runs an agent root through its own lifecycle, and any other node root through NodeRunner. */ + private static Flowable runRoot( + InvocationContext context, Function> processEvent) { + Node root = context.node().orElse(context.agent()); + if (root instanceof BaseAgent agent) { + return agent.runAsync(context).concatMapSingle(processEvent::apply); + } + return NodeRunner.runRoot(root, context, processEvent); + } + private Completable compactEvents(Session session) { return Optional.ofNullable(eventsCompactionConfig) .filter(EventsCompactionConfig::hasSlidingWindowCompactionConfig) @@ -941,8 +1004,14 @@ private Flowable runResumableFromSession( return Flowable.defer( () -> { // Rejected before the invocation-id fallback; defer turns the throw into an error signal. + if (node != null) { + checkNoTextWithFunctionResponses(newMessage); + } if (newMessage != null) { - validateResumeMessage(session, newMessage); + validateResumeMessage( + session, + newMessage, + /* requireFirstId= */ node == null || providedInvocationId != null); } String resolvedInvocationId = resolveInvocationId(session, newMessage, providedInvocationId); @@ -959,6 +1028,11 @@ private Flowable runResumableFromSession( return Flowable.error( new IllegalArgumentException("Session has no events to resume.")); } + if (node != null) { + // Without replay, resuming would run the node again and repeat its side effects. + return Flowable.error( + new UnsupportedOperationException(NODE_ROOT_RESUME_UNSUPPORTED)); + } return resumeCore(session, resolvedInvocationId, newMessage, runConfig, stateDelta); }); } @@ -1227,16 +1301,55 @@ private static Optional matchingFunctionCallEvent(Session session, Conten return Optional.empty(); } + /** + * Rejects a message for a node root if the message mixes text with function responses that carry + * an id: the responses resume an existing invocation, while the text starts a new one. Python's + * node path rejects such a message too; its agent path accepts it. + * + * @throws IllegalArgumentException if {@code message} mixes them. + */ + private static void checkNoTextWithFunctionResponses(@Nullable Content message) { + if (message == null || functionResponseIds(message).stream().allMatch(String::isEmpty)) { + return; + } + boolean hasText = + message.parts().stream() + .flatMap(List::stream) + .map(Part::text) + .flatMap(Optional::stream) + .anyMatch(text -> !text.isEmpty()); + checkArgument( + !hasText, + "Message cannot contain both function responses and text. Function responses resume an" + + " existing invocation while text starts a new one."); + } + + /** + * For a node root, rejects a message whose function responses answer calls from an earlier + * invocation: without replay, running the node again would repeat its side effects. As in Python + * and Kotlin, responses without an id are ignored, so they start a new invocation. + * + * @throws IllegalArgumentException if a response id matches no call, or the ids span invocations. + * @throws UnsupportedOperationException if the responses answer an earlier invocation. + */ + private static void checkNoAnswerToAnEarlierInvocation(Session session, Content message) { + validateResumeMessage(session, message, /* requireFirstId= */ false); + if (matchingFunctionCallEvent(session, message).isPresent()) { + throw new UnsupportedOperationException(NODE_ROOT_RESUME_UNSUPPORTED); + } + } + /** * Rejects a resume message no invocation can be resumed from: one whose first function response - * carries no id, one whose response ids match no call in session history (it would feed the model - * an orphan response), and one whose responses span more than one invocation. Counts rather than - * ids are reported, since the ids come from the caller. Text alongside a function response is - * accepted, as on Python's agent path. + * carries no id when {@code requireFirstId} is set, one whose response ids match no call in + * session history (it would feed the model an orphan response), and one whose responses span more + * than one invocation. Counts rather than ids are reported, since the ids come from the caller. + * Text alongside a function response is accepted, as on Python's agent path. * * @throws IllegalArgumentException if {@code newMessage} is not resumable. */ - private static void validateResumeMessage(Session session, Content newMessage) { + private static void validateResumeMessage( + Session session, Content newMessage, boolean requireFirstId) { ImmutableList responses = newMessage.parts().stream() .flatMap(List::stream) @@ -1249,7 +1362,7 @@ private static void validateResumeMessage(Session session, Content newMessage) { // Python checks the first response only, and raises rather than opening a new invocation that // would carry an orphan response. checkArgument( - responses.get(0).id().isPresent(), + !requireFirstId || responses.get(0).id().isPresent(), "A function response id is required to resume an invocation."); Set unmatched = functionResponseIds(newMessage); Set invocationIds = new HashSet<>(); @@ -1380,7 +1493,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 +1551,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 +1630,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..3c191ff67 --- /dev/null +++ b/core/src/main/java/com/google/adk/workflow/NodeExecutionState.java @@ -0,0 +1,269 @@ +/* + * 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.HashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +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 final ImmutableList outputFor; + private final Map childRunCounts = new HashMap<>(); + 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.outputFor = + builder.useAsOutput && parentState != null + ? ImmutableList.builder().add(nodePath).addAll(parentState.outputFor).build() + : ImmutableList.of(nodePath); + 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 the paths of the nodes that use this run's output as their output, starting with this + * run's path. + */ + ImmutableList outputFor() { + return outputFor; + } + + /** Returns the next run ID for a child node named {@code name}, counting from {@code "1"}. */ + String nextChildRunId(String name) { + return String.valueOf(childRunCounts.merge(name, 1, Integer::sum)); + } + + /** 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 boolean useAsOutput; + + 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; + } + + /** Sets whether this run's output is also its parent's output. */ + @CanIgnoreReturnValue + Builder useAsOutput(boolean useAsOutput) { + this.useAsOutput = useAsOutput; + 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..e2c6ae389 --- /dev/null +++ b/core/src/main/java/com/google/adk/workflow/NodeRunner.java @@ -0,0 +1,501 @@ +/* + * 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.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.Queue; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.function.Function; +import org.jspecify.annotations.Nullable; + +/** + * The runner for one workflow node. It stamps every emitted event with its author, node path, and + * branch before forwarding it, and records the node's output, routes, interrupts, and failure on + * the node state of the returned context. {@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. When {@code nodeFailure} is not null, {@code event} is the error event of a node + * that failed with {@code nodeFailure}. The Completable must not fail; a consumer that cannot + * process an event stops the run by canceling it. + */ + Completable send(Event event, @Nullable Exception nodeFailure); + } + + private final Node node; + private final InvocationContext invocationContext; + private final EventSink eventSink; + private final @Nullable Context parent; + private final String runId; + private final boolean useAsOutput; + 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.useAsOutput = builder.useAsOutput; + 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. If processing an event fails with an exception, the failures reported by + * unprocessed error events are attached to it as suppressed. + */ + @Experimental + public static Flowable runRoot( + Node node, + InvocationContext context, + Function> processEvent) { + checkNotNull(node); + checkNotNull(context); + checkNotNull(processEvent); + return Flowable.defer( + () -> { + FlowableProcessor queue = + UnicastProcessor.create().toSerialized(); + Queue unprocessedFailures = new ConcurrentLinkedQueue<>(); + EventSink sink = + (event, nodeFailure) -> { + CompletableSubject processed = CompletableSubject.create(); + if (nodeFailure != null) { + unprocessedFailures.add(nodeFailure); + } + queue.onNext(new QueuedEvent(event, processed, nodeFailure)); + return event.partial().orElse(false) ? Completable.complete() : processed; + }; + Completable nodeRun = + builder() + .node(node) + .invocationContext(context) + .eventSink(sink) + .build() + .run(context.userContent().orElse(null)) + .doOnSuccess( + nodeContext -> + nodeContext + .requireNodeState() + .failure() + .ifPresentOrElse(queue::onError, queue::onComplete)) + .doOnError(queue::onError) + .ignoreElement() + .onErrorComplete(); + return queue + .concatMap( + queued -> + processEvent + .apply(queued.event()) + .doOnError( + error -> { + // As in Kotlin, keep the failures of unprocessed error events. + if (error instanceof Exception) { + unprocessedFailures.forEach(error::addSuppressed); + } + }) + .toFlowable() + // The node proceeds only after the caller receives the processed event. + .doAfterNext( + unused -> { + if (queued.nodeFailure() != null) { + unprocessedFailures.remove(queued.nodeFailure()); + } + queued.processed().onComplete(); + })) + .mergeWith(nodeRun); + }); + } + + /** + * Runs the node on {@code nodeInput} and returns its context when the node finishes. An exception + * from the node is recorded on the context's node state rather than signaled, while an {@code + * 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) + .useAsOutput(useAsOutput) + .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, /* nodeFailure= */ null) + .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, /* nodeFailure= */ null) + .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, failure).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 ? nodeState.outputFor() : 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 and, for an error event, with the failure it reports. + */ + private record QueuedEvent( + Event event, CompletableSubject processed, @Nullable Exception nodeFailure) {} + + /** Appends a {@code name@runId} segment to {@code base}, joined by {@code separator}. */ + static String appendSegment(@Nullable String base, String name, String runId, char separator) { + String segment = name + "@" + runId; + 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 useAsOutput; + 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's output is also its parent's output, as when a terminal node produces + * the output of its workflow. + */ + @CanIgnoreReturnValue + Builder useAsOutput(boolean useAsOutput) { + this.useAsOutput = useAsOutput; + 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/main/java/com/google/adk/workflow/NodeState.java b/core/src/main/java/com/google/adk/workflow/NodeState.java new file mode 100644 index 000000000..ea6fb867f --- /dev/null +++ b/core/src/main/java/com/google/adk/workflow/NodeState.java @@ -0,0 +1,35 @@ +/* + * 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 com.google.adk.annotations.Experimental; +import com.google.common.collect.ImmutableList; + +/** A node's progress within one run of a workflow graph. */ +@Experimental +final class NodeState { + + /** Where the node stands in the current run. */ + NodeStatus status; + + /** The IDs of the interrupts the node is waiting on, empty unless it is waiting. */ + ImmutableList interrupts = ImmutableList.of(); + + NodeState(NodeStatus status) { + this.status = status; + } +} diff --git a/core/src/main/java/com/google/adk/workflow/NodeStatus.java b/core/src/main/java/com/google/adk/workflow/NodeStatus.java new file mode 100644 index 000000000..f8cf43188 --- /dev/null +++ b/core/src/main/java/com/google/adk/workflow/NodeStatus.java @@ -0,0 +1,32 @@ +/* + * 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 com.google.adk.annotations.Experimental; + +/** The status of a node within one run of a workflow graph. */ +@Experimental +enum NodeStatus { + /** Running now. */ + RUNNING, + + /** Finished; a node that waits for output finishes only once it produces an output or routes. */ + COMPLETED, + + /** Paused, either awaiting a user's answers or awaiting another trigger. */ + WAITING +} diff --git a/core/src/main/java/com/google/adk/workflow/Scheduler.java b/core/src/main/java/com/google/adk/workflow/Scheduler.java new file mode 100644 index 000000000..eb2dbe03c --- /dev/null +++ b/core/src/main/java/com/google/adk/workflow/Scheduler.java @@ -0,0 +1,334 @@ +/* + * 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.Strings.isNullOrEmpty; +import static com.google.common.collect.ImmutableList.toImmutableList; + +import com.google.adk.agents.Context; +import com.google.adk.annotations.Experimental; +import com.google.common.base.Splitter; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; +import io.reactivex.rxjava3.core.Completable; +import io.reactivex.rxjava3.disposables.CompositeDisposable; +import io.reactivex.rxjava3.disposables.Disposable; +import io.reactivex.rxjava3.processors.FlowableProcessor; +import io.reactivex.rxjava3.processors.UnicastProcessor; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Queue; +import java.util.Set; +import java.util.concurrent.ConcurrentLinkedQueue; +import org.jspecify.annotations.Nullable; + +/** One run of one workflow's scheduling loop. Discarded when the run ends. */ +@Experimental +final class Scheduler { + + private final Workflow workflow; + private final Graph graph; + private final Context context; + private final NodeExecutionState nodeState; + private final Map nodeStates = new HashMap<>(); + private final Map nodeOutputs = new HashMap<>(); + // The branch each completed node ran on, so a join can re-merge parallel branches. + private final Map nodeBranches = new HashMap<>(); + private final TriggerQueue triggerQueue = new TriggerQueue(); + private final Set interruptIds = new LinkedHashSet<>(); + + Scheduler(Workflow workflow, Graph graph, Context context) { + this.workflow = workflow; + this.graph = graph; + this.context = context; + this.nodeState = context.requireNodeState(); + } + + /** Runs the graph from its start node, which hands {@code nodeInput} to its successors. */ + Completable run(@Nullable Object nodeInput) { + return Completable.defer( + () -> { + enqueueStartTriggers(nodeInput); + return new Loop().drive(); + }); + } + + private void enqueueStartTriggers(@Nullable Object nodeInput) { + ImmutableList startEdges = + graph.edges().stream() + .filter(edge -> edge.from().name().equals(Start.NAME)) + .collect(toImmutableList()); + // A start node with more than one successor puts each on its own sub-branch. + boolean useSubBranch = startEdges.size() > 1; + for (Edge edge : startEdges) { + if (edge.to().requiresAllPredecessors()) { + // The node must still wait for its other predecessors, so record the start node's input. + nodeOutputs.put(Start.NAME, nodeInput); + nodeBranches.put(Start.NAME, context.invocationContext().branch().orElse(null)); + enqueueBarrierTrigger(edge.to().name()); + continue; + } + triggerQueue.enqueue( + edge.to().name(), new Trigger(nodeInput, useSubBranch, /* branch= */ null)); + } + } + + /** + * The scheduling loop. Node runs report their completions through a serialized signal that the + * loop handles one batch at a time, so its state needs no locks. It runs until no node is running + * or can start, or until a node fails. + */ + private final class Loop { + private final Map running = new LinkedHashMap<>(); + private final CompositeDisposable childRuns = new CompositeDisposable(); + private final Queue completions = new ConcurrentLinkedQueue<>(); + private final FlowableProcessor wakeUps = + UnicastProcessor.create().toSerialized(); + private boolean failed; + + Completable drive() { + return Completable.defer( + () -> { + scheduleReady(); + stopIfIdle(); + return wakeUps.doOnNext(unused -> handleCompletions()).ignoreElements(); + }) + .andThen(Completable.defer(() -> failed ? Completable.complete() : finalizeRun())) + .doFinally(childRuns::dispose); + } + + /** Handles the node runs that finished since the last wake-up. */ + private void handleCompletions() { + // A failure ends the run, so a completion queued behind it must not start successors. + if (failed) { + return; + } + List batch = new ArrayList<>(); + for (Completion completion; (completion = completions.poll()) != null; ) { + batch.add(completion); + childRuns.delete(running.remove(completion.nodeName())); + } + if (batch.isEmpty()) { + return; + } + // Act on a failure first, or the successors of a node that finished beside it would start. + if (handleFailures(batch)) { + return; + } + batch.forEach(completion -> handleCompletion(completion.nodeName(), completion.context())); + scheduleReady(); + stopIfIdle(); + } + + /** + * Fails the run with the first failed node in {@code batch}, canceling the nodes still running, + * and returns whether a node failed. + */ + private boolean handleFailures(List batch) { + for (Completion completion : batch) { + Optional failure = completion.context().requireNodeState().failure(); + if (failure.isPresent()) { + nodeState.failure(failure.get()); + failed = true; + childRuns.dispose(); + wakeUps.onComplete(); + return true; + } + } + return false; + } + + private void stopIfIdle() { + if (running.isEmpty()) { + wakeUps.onComplete(); + } + } + + private void scheduleReady() { + Optional maxConcurrency = workflow.maxConcurrency(); + for (String nodeName : triggerQueue.queuedNodeNames()) { + if (running.containsKey(nodeName)) { + continue; + } + // A node waiting on a user's answers keeps its triggers until it is resumed. + NodeState state = nodeStates.get(nodeName); + if (state != null && state.status == NodeStatus.WAITING && !state.interrupts.isEmpty()) { + continue; + } + if (maxConcurrency.isPresent() && running.size() >= maxConcurrency.get()) { + break; + } + triggerQueue.poll(nodeName).ifPresent(trigger -> start(nodeName, trigger)); + } + } + + private void start(String nodeName, Trigger trigger) { + String runId = nodeState.nextChildRunId(nodeName); + nodeStates.put(nodeName, new NodeState(NodeStatus.RUNNING)); + Disposable run = + NodeRunner.builder() + .node(graph.node(nodeName)) + .parent(context) + .runId(runId) + .useAsOutput(graph.terminalNodeNames().contains(nodeName)) + .useSubBranch(trigger.useSubBranch()) + .overrideBranch(trigger.branch()) + .build() + .run(trigger.input()) + .subscribe( + childContext -> { + completions.add(new Completion(nodeName, childContext)); + wakeUps.onNext(true); + }, + // A node's exception fails only that node, so this is an Error. + wakeUps::onError); + childRuns.add(run); + running.put(nodeName, run); + } + } + + private void handleCompletion(String nodeName, Context childContext) { + NodeState state = nodeStates.get(nodeName); + NodeExecutionState childState = childContext.requireNodeState(); + ImmutableSet childInterruptIds = childState.interruptIds(); + if (!childInterruptIds.isEmpty()) { + state.status = NodeStatus.WAITING; + state.interrupts = childInterruptIds.asList(); + interruptIds.addAll(childInterruptIds); + return; + } + if (graph.node(nodeName).waitForOutput() + && childState.output().isEmpty() + && childState.selectedRoutes().isEmpty()) { + state.status = NodeStatus.WAITING; + return; + } + state.status = NodeStatus.COMPLETED; + nodeBranches.put(nodeName, childContext.invocationContext().branch().orElse(null)); + childState.output().ifPresent(output -> nodeOutputs.put(nodeName, output)); + enqueueSuccessorTriggers(nodeName, childContext); + } + + private void enqueueSuccessorTriggers(String nodeName, Context childContext) { + NodeExecutionState childState = childContext.requireNodeState(); + ImmutableList successors = + graph.nodesTriggeredBy(nodeName, childState.selectedRoutes().orElse(ImmutableList.of())); + // Several successors each run on a sub-branch, while a single one keeps the node's branch. + boolean useSubBranch = successors.size() > 1; + String branch = childContext.invocationContext().branch().orElse(null); + for (String successor : successors) { + if (graph.node(successor).requiresAllPredecessors()) { + enqueueBarrierTrigger(successor); + } else { + triggerQueue.enqueue( + successor, new Trigger(childState.output().orElse(null), useSubBranch, branch)); + } + } + } + + /** + * Queues a trigger for {@code target}, a node that waits for all its predecessors, once every + * predecessor has completed. The trigger hands it their outputs keyed by node name. + */ + private void enqueueBarrierTrigger(String target) { + ImmutableSet predecessors = graph.predecessorsOf(target); + // A predecessor that a loop re-triggered is still marked completed, so check its queue as well. + ImmutableSet queued = ImmutableSet.copyOf(triggerQueue.queuedNodeNames()); + for (String predecessor : predecessors) { + NodeState state = nodeStates.get(predecessor); + boolean done = + predecessor.equals(Start.NAME) + || (state != null + && state.status == NodeStatus.COMPLETED + && !queued.contains(predecessor)); + if (!done) { + return; + } + } + Map outputs = new LinkedHashMap<>(); + List<@Nullable String> branches = new ArrayList<>(); + for (String predecessor : predecessors) { + outputs.put(predecessor, nodeOutputs.get(predecessor)); + branches.add(nodeBranches.get(predecessor)); + } + // The node re-merges onto the branch that its predecessors forked from. + triggerQueue.enqueue( + target, + new Trigger( + Collections.unmodifiableMap(outputs), + /* useSubBranch= */ false, + commonBranchPrefix(branches))); + } + + /** + * Surfaces the interrupts of the nodes still waiting on the workflow, or, when there are none, + * records the output of the sole terminal node that produced one as the workflow's output. + */ + private Completable finalizeRun() { + if (!interruptIds.isEmpty()) { + nodeState.addInterruptIds(interruptIds); + return Completable.complete(); + } + ImmutableList terminalOutputs = + graph.terminalNodeNames().stream() + .filter(nodeOutputs::containsKey) + .collect(toImmutableList()); + if (terminalOutputs.size() > 1) { + return Completable.error( + new WorkflowConfigurationException( + String.format( + "Workflow '%s': multiple terminal nodes produced output (%d). A workflow must" + + " have at most one terminal output.", + workflow.name(), terminalOutputs.size()))); + } + if (terminalOutputs.size() == 1) { + nodeState.produceOutput(nodeOutputs.get(terminalOutputs.get(0))); + // The terminal node's event already carried the output for the workflow as well. + nodeState.markOutputEmitted(); + } + return Completable.complete(); + } + + /** Returns the longest segment prefix that the dot-separated {@code branches} share, if any. */ + private static @Nullable String commonBranchPrefix(List<@Nullable String> branches) { + List> segmented = new ArrayList<>(); + for (String branch : branches) { + segmented.add( + isNullOrEmpty(branch) ? ImmutableList.of() : Splitter.on('.').splitToList(branch)); + } + int shortest = segmented.stream().mapToInt(List::size).min().getAsInt(); + List common = new ArrayList<>(); + for (int i = 0; i < shortest; i++) { + String segment = segmented.get(0).get(i); + int index = i; + if (segmented.stream().anyMatch(segments -> !segments.get(index).equals(segment))) { + break; + } + common.add(segment); + } + return common.isEmpty() ? null : String.join(".", common); + } + + /** A node run that finished, with its context. */ + private record Completion(String nodeName, Context context) {} +} diff --git a/core/src/main/java/com/google/adk/workflow/Trigger.java b/core/src/main/java/com/google/adk/workflow/Trigger.java new file mode 100644 index 000000000..f619c1e71 --- /dev/null +++ b/core/src/main/java/com/google/adk/workflow/Trigger.java @@ -0,0 +1,35 @@ +/* + * 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 com.google.adk.annotations.Experimental; +import org.jspecify.annotations.Nullable; + +/** + * A queued reason to run a node. + * + * @param input the node's input: the workflow's input for a successor of the start node, the + * triggering node's output, or, for a node that waits for all its predecessors, their outputs + * keyed by node name + * @param useSubBranch whether the node runs on its own sub-branch, as when the triggering node had + * more than one successor + * @param branch the branch the node inherits instead of the workflow's branch: the triggering + * node's branch (which a sub-branch extends) or the branch that a join's predecessors forked + * from + */ +@Experimental +record Trigger(@Nullable Object input, boolean useSubBranch, @Nullable String branch) {} diff --git a/core/src/main/java/com/google/adk/workflow/TriggerQueue.java b/core/src/main/java/com/google/adk/workflow/TriggerQueue.java new file mode 100644 index 000000000..ce286c599 --- /dev/null +++ b/core/src/main/java/com/google/adk/workflow/TriggerQueue.java @@ -0,0 +1,58 @@ +/* + * 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 com.google.adk.annotations.Experimental; +import com.google.common.collect.ImmutableList; +import java.util.ArrayDeque; +import java.util.Deque; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.Optional; + +/** + * The pending triggers of each node in a workflow run. Each node keeps its own queue, so the + * triggers of parallel branches or loops are served in arrival order without affecting other nodes. + */ +@Experimental +final class TriggerQueue { + + private final Map> queues = new LinkedHashMap<>(); + + /** Returns the names of the nodes with at least one pending trigger, in first-queued order. */ + ImmutableList queuedNodeNames() { + return ImmutableList.copyOf(queues.keySet()); + } + + /** Queues {@code trigger} for the node named {@code nodeName}. */ + void enqueue(String nodeName, Trigger trigger) { + queues.computeIfAbsent(nodeName, unused -> new ArrayDeque<>()).addLast(trigger); + } + + /** Removes and returns the next pending trigger of the node named {@code nodeName}, if any. */ + Optional poll(String nodeName) { + Deque queue = queues.get(nodeName); + if (queue == null) { + return Optional.empty(); + } + Trigger trigger = queue.removeFirst(); + if (queue.isEmpty()) { + queues.remove(nodeName); + } + return Optional.of(trigger); + } +} diff --git a/core/src/main/java/com/google/adk/workflow/Workflow.java b/core/src/main/java/com/google/adk/workflow/Workflow.java index 4d2dc00db..de88384d5 100644 --- a/core/src/main/java/com/google/adk/workflow/Workflow.java +++ b/core/src/main/java/com/google/adk/workflow/Workflow.java @@ -25,6 +25,7 @@ import com.google.common.collect.ImmutableList; import com.google.errorprone.annotations.CanIgnoreReturnValue; import com.google.genai.types.Schema; +import io.reactivex.rxjava3.core.Completable; import io.reactivex.rxjava3.core.Flowable; import java.util.List; import java.util.Optional; @@ -78,10 +79,23 @@ Optional graph() { return Optional.ofNullable(graph); } + /** + * Runs the graph from its start node, which hands {@code nodeInput} to its successors. The output + * of the sole terminal node that produces one becomes the workflow's output. + */ @Override public Flowable runNode(Context context, @Nullable Object nodeInput) { - return Flowable.error( - new UnsupportedOperationException("Running a workflow is not supported yet.")); + Graph runGraph = graph; + if (runGraph == null) { + return Flowable.empty(); + } + return Completable.defer( + () -> { + // Inner nodes inherit the workflow as event author; agent nodes keep their own. + context.requireNodeState().eventAuthor(name()); + return new Scheduler(this, runGraph, context).run(nodeInput); + }) + .toFlowable(); } /** diff --git a/core/src/test/java/com/google/adk/workflow/StubNode.java b/core/src/main/java/com/google/adk/workflow/WorkflowConfigurationException.java similarity index 63% rename from core/src/test/java/com/google/adk/workflow/StubNode.java rename to core/src/main/java/com/google/adk/workflow/WorkflowConfigurationException.java index 80a7e0d82..b75e21b20 100644 --- a/core/src/test/java/com/google/adk/workflow/StubNode.java +++ b/core/src/main/java/com/google/adk/workflow/WorkflowConfigurationException.java @@ -16,14 +16,13 @@ package com.google.adk.workflow; -import com.google.adk.agents.Context; -import io.reactivex.rxjava3.core.Flowable; -import org.jspecify.annotations.Nullable; +import com.google.adk.annotations.Experimental; -/** A node that emits nothing, for tests that only need a named node in a graph. */ -record StubNode(String name) implements Node { - @Override - public Flowable runNode(Context context, @Nullable Object nodeInput) { - return Flowable.empty(); +/** Thrown when a workflow or one of its nodes is declared in a way that cannot run. */ +@Experimental +public final class WorkflowConfigurationException extends IllegalArgumentException { + + public WorkflowConfigurationException(String message) { + super(message); } } 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..4d803edbe 100644 --- a/core/src/test/java/com/google/adk/agents/BaseAgentTest.java +++ b/core/src/test/java/com/google/adk/agents/BaseAgentTest.java @@ -25,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 @@ -480,6 +468,26 @@ public void runLive_invokesRunLiveImpl() { assertThat(runLiveCallback.wasCalled()).isTrue(); } + @Test + public void runNode_withoutANodeContext_throwsBeforeRunningTheAgent() { + var beforeCallback = TestCallback.returningEmpty(); + var runAsyncImpl = TestCallback.returningEmpty(); + TestBaseAgent agent = + new TestBaseAgent( + TEST_AGENT_NAME, + TEST_AGENT_DESCRIPTION, + ImmutableList.of(beforeCallback.asBeforeAgentCallback()), + /* afterAgentCallbacks= */ ImmutableList.of(), + runAsyncImpl.asRunAsyncImplSupplier(Content.fromParts(Part.fromText("main_output")))); + CallbackContext context = + new CallbackContext(TestUtils.createInvocationContext(agent), /* eventActions= */ null); + + assertThrows(IllegalStateException.class, () -> agent.runNode(context, /* nodeInput= */ null)); + + assertThat(beforeCallback.wasCalled()).isFalse(); + assertThat(runAsyncImpl.wasCalled()).isFalse(); + } + @Test public void runLive_beforeAgentCallbackReturnsContent_endsInvocationAndSkipsRunLiveImplAndAfterCallback() { 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..598b1739a 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_throwsIllegalStateException() { + CallbackContext context = new CallbackContext(newInvocationContext(), null); + + assertThrows(IllegalStateException.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..7e60469b3 --- /dev/null +++ b/core/src/test/java/com/google/adk/runner/RunnerRootNodeTest.java @@ -0,0 +1,686 @@ +/* + * 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.Edge; +import com.google.adk.workflow.Node; +import com.google.adk.workflow.NodeInfo; +import com.google.adk.workflow.Start; +import com.google.adk.workflow.Workflow; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.genai.types.Content; +import com.google.genai.types.FinishReason; +import com.google.genai.types.FunctionCall; +import com.google.genai.types.FunctionResponse; +import com.google.genai.types.Part; +import io.reactivex.rxjava3.core.Completable; +import io.reactivex.rxjava3.core.Flowable; +import io.reactivex.rxjava3.core.Maybe; +import io.reactivex.rxjava3.subscribers.TestSubscriber; +import java.util.ArrayList; +import java.util.List; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests for running a workflow node, rather than an agent, at the root of an invocation. */ +@RunWith(JUnit4.class) +public final class RunnerRootNodeTest { + + private static final String APP_NAME = "node_app"; + private static final String USER_ID = "user"; + private static final Content USER_MESSAGE = Content.fromParts(Part.fromText("hello")); + private static final FunctionResponse IDENTIFIED_RESPONSE = + FunctionResponse.builder().id("fc-1").name("tool").response(ImmutableMap.of("v", 1)).build(); + private static final FunctionResponse UNIDENTIFIED_RESPONSE = + FunctionResponse.builder().name("tool").response(ImmutableMap.of("v", 1)).build(); + + @Test + public void node_ofARunnerForANodeApp_isThatNode() { + Node node = TestNode.of("solo"); + + Runner runner = runner(App.builder().name(APP_NAME).rootNode(node).build()); + + assertThat(runner.node()).hasValue(node); + } + + @Test + public void build_withANodeViewAgentAsTheAgent_throwsIllegalArgumentException() { + Runner.Builder builder = + Runner.builder().agent(new NodeViewAgent(TestNode.of("solo"))).appName(APP_NAME); + + assertThrows(IllegalArgumentException.class, builder::build); + } + + @Test + @SuppressWarnings("deprecation") // The deprecated constructors must reject it as well. + public void constructor_withANodeViewAgent_throwsIllegalArgumentException() { + NodeViewAgent agent = new NodeViewAgent(TestNode.of("solo")); + + assertThrows( + IllegalArgumentException.class, + () -> + new Runner( + agent, APP_NAME, new InMemoryArtifactService(), new InMemorySessionService())); + } + + @Test + public void runAsync_nodeRoot_runsTheNodeOnTheUserMessage() { + List inputs = new ArrayList<>(); + Node node = + new TestNode( + "solo", + (context, input) -> { + inputs.add(input); + return Flowable.empty(); + }); + Runner runner = runner(node); + Session session = newSession(runner); + + runner.runAsync(USER_ID, session.id(), USER_MESSAGE).blockingSubscribe(); + + assertThat(inputs).containsExactly(USER_MESSAGE); + } + + @Test + public void runAsync_nodeRoot_persistsItsStampedEventAndState() { + Node node = + new TestNode( + "solo", + (context, input) -> { + context.state().put("durable", "saved"); + return Flowable.just("echoed"); + }); + Runner runner = runner(node); + Session session = newSession(runner); + + List events = + runner.runAsync(USER_ID, session.id(), USER_MESSAGE).toList().blockingGet(); + + Event event = getOnlyElement(events); + Session saved = reload(runner, session); + assertThat(event.author()).isEqualTo("solo"); + assertThat(event.nodeInfo().map(NodeInfo::path)).hasValue("solo@1"); + assertThat(event.output()).hasValue("echoed"); + assertThat(getLast(saved.immutableEvents()).id()).isEqualTo(event.id()); + assertThat(saved.state()).containsEntry("durable", "saved"); + } + + @Test + public void runAsync_workflowRoot_runsItsGraphAndPersistsEveryEvent() { + Node first = new TestNode("first", (context, input) -> Flowable.just("one")); + Node second = new TestNode("second", (context, input) -> Flowable.just("two")); + Workflow workflow = + Workflow.builder() + .name("flow") + .edges(new Edge(Start.INSTANCE, first), new Edge(first, second)) + .build(); + Runner runner = runner(workflow); + Session session = newSession(runner); + + List events = + runner.runAsync(USER_ID, session.id(), USER_MESSAGE).toList().blockingGet(); + + Session saved = reload(runner, session); + assertThat( + events.stream() + .map(event -> event.nodeInfo().map(NodeInfo::path)) + .collect(toImmutableList())) + .containsExactly(Optional.of("flow@1/first@1"), Optional.of("flow@1/second@1")) + .inOrder(); + assertThat(events.stream().map(Event::author).collect(toImmutableList())) + .containsExactly("flow", "flow"); + assertThat(getLast(saved.immutableEvents()).id()).isEqualTo(getLast(events).id()); + } + + @Test + public void close_workflowRoot_closesTheAgentsInItsGraphAndNestedWorkflows() { + List closed = new ArrayList<>(); + BaseAgent outerAgent = closeRecordingAgent("outer_agent", closed); + BaseAgent innerAgent = closeRecordingAgent("inner_agent", closed); + Workflow inner = + Workflow.builder().name("inner").edges(new Edge(Start.INSTANCE, innerAgent)).build(); + Workflow workflow = + Workflow.builder() + .name("flow") + .edges(new Edge(Start.INSTANCE, outerAgent), new Edge(outerAgent, inner)) + .build(); + + runner(workflow).close().blockingAwait(); + + assertThat(closed).containsExactly("outer_agent", "inner_agent"); + } + + @Test + public void runAsync_failingNodeRoot_persistsItsErrorEventBeforeFailing() { + IllegalStateException failure = new IllegalStateException("node failed"); + Runner runner = runner(failingNode(failure)); + Session session = newSession(runner); + + TestSubscriber subscriber = runner.runAsync(USER_ID, session.id(), USER_MESSAGE).test(); + subscriber.awaitDone(5, SECONDS); + + subscriber.assertError(failure); + Event errorEvent = getOnlyElement(subscriber.values()); + Session saved = reload(runner, session); + assertThat(errorEvent.errorCode()).hasValue(new FinishReason("IllegalStateException")); + assertThat(errorEvent.nodeInfo().map(NodeInfo::path)).hasValue("boom@1"); + assertThat(getLast(saved.immutableEvents()).id()).isEqualTo(errorEvent.id()); + assertThat(saved.state()).containsEntry("partial_step", "done"); + } + + @Test + public void runAsync_failingNodeRoot_reportsTheErrorToPluginsWithoutAfterRun() { + IllegalStateException failure = new IllegalStateException("node failed"); + RecordingPlugin plugin = new RecordingPlugin(); + Runner runner = runner(failingNode(failure), plugin); + Session session = newSession(runner); + + TestSubscriber subscriber = runner.runAsync(USER_ID, session.id(), USER_MESSAGE).test(); + subscriber.awaitDone(5, SECONDS); + + Event errorEvent = getOnlyElement(subscriber.values()); + assertThat(ids(plugin.events)).containsExactly(errorEvent.id()); + assertThat(plugin.runErrors).containsExactly(failure); + assertThat(plugin.afterRunCount).isEqualTo(0); + } + + @Test + public void runAsync_onEventCallbackFailure_cancelsTheNodeRootAndFails() { + IllegalStateException callbackFailure = new IllegalStateException("onEvent failed"); + AtomicBoolean cancelled = new AtomicBoolean(); + Node node = + new TestNode( + "emitter", + (context, input) -> + Flowable.just(contentEvent("first")) + .concatWith(Flowable.never()) + .doOnCancel(() -> cancelled.set(true))); + Plugin failingPlugin = + new BasePlugin("failing") { + @Override + public Maybe onEventCallback(InvocationContext invocationContext, Event event) { + return Maybe.error(callbackFailure); + } + }; + Runner runner = runner(node, failingPlugin); + Session session = newSession(runner); + + TestSubscriber subscriber = runner.runAsync(USER_ID, session.id(), USER_MESSAGE).test(); + subscriber.awaitDone(5, SECONDS); + + subscriber.assertError(callbackFailure); + subscriber.assertNoValues(); + assertThat(cancelled.get()).isTrue(); + } + + @Test + public void runAsync_beforeRunCallbackContent_skipsTheNodeRootAndRunsAfterRun() { + Content halted = Content.fromParts(Part.fromText("halted by plugin")); + AtomicBoolean nodeRan = new AtomicBoolean(); + Node node = + new TestNode( + "should_not_run", + (context, input) -> { + nodeRan.set(true); + return Flowable.just("unreachable"); + }); + RecordingPlugin plugin = + new RecordingPlugin() { + @Override + public Maybe beforeRunCallback(InvocationContext invocationContext) { + return Maybe.just(halted); + } + }; + Runner runner = runner(node, plugin); + Session session = newSession(runner); + + List events = + runner.runAsync(USER_ID, session.id(), USER_MESSAGE).toList().blockingGet(); + + assertThat(nodeRan.get()).isFalse(); + assertThat(getOnlyElement(events).content()).hasValue(halted); + assertThat(plugin.afterRunCount).isEqualTo(1); + } + + @Test + public void runAsync_onUserMessageCallback_replacesTheNodeRootsInput() { + Content rewritten = Content.fromParts(Part.fromText("modified text")); + List inputs = new ArrayList<>(); + Node node = + new TestNode( + "recorder", + (context, input) -> { + inputs.add(input); + return Flowable.empty(); + }); + Plugin rewritingPlugin = + new BasePlugin("rewriting") { + @Override + public Maybe onUserMessageCallback( + InvocationContext invocationContext, Content userMessage) { + return Maybe.just(rewritten); + } + }; + Runner runner = runner(node, rewritingPlugin); + Session session = newSession(runner); + + runner.runAsync(USER_ID, session.id(), USER_MESSAGE).blockingSubscribe(); + + assertThat(inputs).containsExactly(rewritten); + } + + @Test + public void runAsync_partialEventOfANodeRoot_isStreamedButNotPersisted() { + Node node = + new TestNode( + "streamer", + (context, input) -> + Flowable.just( + contentEvent("part").toBuilder().partial(true).build(), contentEvent("final"))); + Runner runner = runner(node); + Session session = newSession(runner); + + List events = + runner.runAsync(USER_ID, session.id(), USER_MESSAGE).toList().blockingGet(); + + assertThat(texts(events)).containsExactly("part", "final").inOrder(); + assertThat( + texts( + reload(runner, session).immutableEvents().stream() + .filter(event -> event.author().equals("streamer")) + .collect(toImmutableList()))) + .containsExactly("final"); + } + + @Test + public void runAsync_takingTheFirstEventOfANodeRoot_cancelsTheNode() { + AtomicBoolean cancelled = new AtomicBoolean(); + Node node = + new TestNode( + "streamer", + (context, input) -> + Flowable.just(contentEvent("first")) + .concatWith(Flowable.never()) + .doOnCancel(() -> cancelled.set(true))); + Runner runner = runner(node); + Session session = newSession(runner); + + TestSubscriber subscriber = + runner.runAsync(USER_ID, session.id(), USER_MESSAGE).take(1).test(); + subscriber.awaitDone(5, SECONDS); + + subscriber.assertComplete(); + assertThat(texts(subscriber.values())).containsExactly("first"); + assertThat(cancelled.get()).isTrue(); + } + + @Test + public void runAsync_resumingANodeRootInvocation_failsWithoutRunningTheNodeAgain() { + AtomicInteger runs = new AtomicInteger(); + Node node = + new TestNode( + "solo", + (context, input) -> { + runs.incrementAndGet(); + return Flowable.just("done"); + }); + Runner runner = resumableRunner(node); + Session session = newSession(runner); + String invocationId = + runner.runAsync(USER_ID, session.id(), USER_MESSAGE).blockingFirst().invocationId(); + + TestSubscriber subscriber = + runner + .runAsync( + USER_ID, + session.id(), + invocationId, + /* newMessage= */ null, + RunConfig.builder().build(), + /* stateDelta= */ null) + .test(); + subscriber.awaitDone(5, SECONDS); + + subscriber.assertError(UnsupportedOperationException.class); + assertThat(runs.get()).isEqualTo(1); + } + + @Test + public void runAsync_textWithAnIdentifiedFunctionResponseToANodeRoot_failsBeforeAnyEvent() { + AtomicInteger runs = new AtomicInteger(); + Runner runner = runner(countingNode(runs)); + Session session = newSession(runner); + Content message = textWithFunctionResponse(IDENTIFIED_RESPONSE); + Flowable events = runner.runAsync(USER_ID, session.id(), message); + + IllegalArgumentException error = + assertThrows(IllegalArgumentException.class, events::blockingSubscribe); + + assertThat(error).hasMessageThat().contains("both function responses and text"); + assertThat(runs.get()).isEqualTo(0); + assertThat(reload(runner, session).immutableEvents()).isEmpty(); + } + + @Test + public void runAsync_textWithAnIdentifiedFunctionResponseToAResumableNodeRoot_fails() { + AtomicInteger runs = new AtomicInteger(); + Runner runner = resumableRunner(countingNode(runs)); + Session session = newSession(runner); + Content message = textWithFunctionResponse(IDENTIFIED_RESPONSE); + Flowable events = runner.runAsync(USER_ID, session.id(), message); + + IllegalArgumentException error = + assertThrows(IllegalArgumentException.class, events::blockingSubscribe); + + assertThat(error).hasMessageThat().contains("both function responses and text"); + assertThat(runs.get()).isEqualTo(0); + } + + @Test + public void runAsync_textWithAFunctionResponseWithoutIdToANodeRoot_runsTheNode() { + AtomicInteger runs = new AtomicInteger(); + Runner runner = runner(countingNode(runs)); + Session session = newSession(runner); + + runner + .runAsync(USER_ID, session.id(), textWithFunctionResponse(UNIDENTIFIED_RESPONSE)) + .blockingSubscribe(); + + assertThat(runs.get()).isEqualTo(1); + } + + @Test + public void runAsync_functionResponseToTheNodeRootsEarlierCall_failsWithoutRunningTheNodeAgain() { + AtomicInteger runs = new AtomicInteger(); + Node node = + new TestNode( + "solo", + (context, input) -> { + runs.incrementAndGet(); + return Flowable.just(functionCallEvent(IDENTIFIED_RESPONSE.id().orElseThrow())); + }); + Runner runner = runner(node); + Session session = newSession(runner); + runner.runAsync(USER_ID, session.id(), USER_MESSAGE).blockingSubscribe(); + int eventsAfterTheFirstRun = reload(runner, session).immutableEvents().size(); + Flowable events = + runner.runAsync(USER_ID, session.id(), functionResponse(IDENTIFIED_RESPONSE)); + + assertThrows(UnsupportedOperationException.class, events::blockingSubscribe); + + assertThat(runs.get()).isEqualTo(1); + assertThat(reload(runner, session).immutableEvents()).hasSize(eventsAfterTheFirstRun); + } + + @Test + public void runAsync_functionResponseWithAnUnknownIdToANodeRoot_failsBeforeAnyEvent() { + AtomicInteger runs = new AtomicInteger(); + Runner runner = runner(countingNode(runs)); + Session session = newSession(runner); + Flowable events = + runner.runAsync(USER_ID, session.id(), functionResponse(IDENTIFIED_RESPONSE)); + + IllegalArgumentException error = + assertThrows(IllegalArgumentException.class, events::blockingSubscribe); + + assertThat(error).hasMessageThat().contains("No matching function call"); + assertThat(runs.get()).isEqualTo(0); + assertThat(reload(runner, session).immutableEvents()).isEmpty(); + } + + @Test + public void runAsync_functionResponseWithoutIdToAResumableNodeRoot_runsTheNode() { + AtomicInteger runs = new AtomicInteger(); + Runner runner = resumableRunner(countingNode(runs)); + Session session = newSession(runner); + + runner + .runAsync(USER_ID, session.id(), functionResponse(UNIDENTIFIED_RESPONSE)) + .blockingSubscribe(); + + assertThat(runs.get()).isEqualTo(1); + } + + @Test + public void runAsync_textWithAFunctionResponseWithoutIdToAResumableNodeRoot_runsTheNode() { + AtomicInteger runs = new AtomicInteger(); + Runner runner = resumableRunner(countingNode(runs)); + Session session = newSession(runner); + + runner + .runAsync(USER_ID, session.id(), textWithFunctionResponse(UNIDENTIFIED_RESPONSE)) + .blockingSubscribe(); + + assertThat(runs.get()).isEqualTo(1); + } + + @Test + public void runAsync_resumingANodeRootWithAFunctionResponseWithoutId_failsBeforeRunningTheNode() { + AtomicInteger runs = new AtomicInteger(); + Runner runner = resumableRunner(countingNode(runs)); + Session session = newSession(runner); + runner.runAsync(USER_ID, session.id(), USER_MESSAGE).blockingSubscribe(); + String invocationId = getLast(reload(runner, session).immutableEvents()).invocationId(); + Flowable events = + runner.runAsync( + USER_ID, + session.id(), + invocationId, + functionResponse(UNIDENTIFIED_RESPONSE), + RunConfig.builder().build(), + /* stateDelta= */ null); + + IllegalArgumentException error = + assertThrows(IllegalArgumentException.class, events::blockingSubscribe); + + assertThat(error).hasMessageThat().contains("function response id is required"); + assertThat(runs.get()).isEqualTo(1); + } + + @Test + public void runAsync_textWithAnIdentifiedFunctionResponseToAnAgentRoot_runsTheAgent() { + TestBaseAgent agent = new TestBaseAgent("solo_agent", "An agent.", null, null, Flowable::empty); + Runner runner = runner(App.builder().name(APP_NAME).rootAgent(agent).build()); + Session session = newSession(runner); + + runner + .runAsync(USER_ID, session.id(), textWithFunctionResponse(IDENTIFIED_RESPONSE)) + .blockingSubscribe(); + + assertThat(agent.getInvocationCount()).isEqualTo(1); + } + + @Test + public void runLive_nodeRoot_failsWithUnsupportedOperationExceptionBeforeAnyAgentCallback() { + AtomicBoolean agentCallbackRan = new AtomicBoolean(); + Plugin agentCallbackRecorder = + new BasePlugin("agent_callback_recorder") { + @Override + public Maybe beforeAgentCallback( + BaseAgent agent, CallbackContext callbackContext) { + agentCallbackRan.set(true); + return Maybe.empty(); + } + }; + Runner runner = runner(TestNode.of("solo"), agentCallbackRecorder); + Session session = newSession(runner); + + TestSubscriber subscriber = + runner.runLive(session, new LiveRequestQueue(), RunConfig.builder().build()).test(); + subscriber.awaitDone(5, SECONDS); + + subscriber.assertError(UnsupportedOperationException.class); + assertThat(agentCallbackRan.get()).isFalse(); + } + + private static Runner runner(Node node, Plugin... plugins) { + return runner(App.builder().name(APP_NAME).rootNode(node).plugins(plugins).build()); + } + + private static Runner runner(App app) { + return Runner.builder().app(app).build(); + } + + private static Runner resumableRunner(Node node) { + return runner( + App.builder() + .name(APP_NAME) + .rootNode(node) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build()); + } + + private static Session newSession(Runner runner) { + return runner.sessionService().createSession(APP_NAME, USER_ID).blockingGet(); + } + + private static Session reload(Runner runner, Session session) { + return runner + .sessionService() + .getSession(APP_NAME, USER_ID, session.id(), Optional.empty()) + .blockingGet(); + } + + private static BaseAgent closeRecordingAgent(String name, List closed) { + return new TestBaseAgent(name, "An agent.", null, null, Flowable::empty) { + @Override + public Completable close() { + closed.add(name()); + return Completable.complete(); + } + }; + } + + private static Node failingNode(Exception failure) { + return new TestNode( + "boom", + (context, input) -> { + context.state().put("partial_step", "done"); + return Flowable.error(failure); + }); + } + + private static Node countingNode(AtomicInteger runs) { + return new TestNode( + "solo", + (context, input) -> { + runs.incrementAndGet(); + return Flowable.empty(); + }); + } + + private static Content functionResponse(FunctionResponse response) { + return Content.builder() + .role("user") + .parts(Part.builder().functionResponse(response).build()) + .build(); + } + + private static Content textWithFunctionResponse(FunctionResponse response) { + return Content.builder() + .role("user") + .parts(Part.fromText("some text"), Part.builder().functionResponse(response).build()) + .build(); + } + + private static Event contentEvent(String text) { + return Event.builder() + .id(Event.generateEventId()) + .author("") + .content(Content.fromParts(Part.fromText(text))) + .build(); + } + + private static Event functionCallEvent(String callId) { + return Event.builder() + .id(Event.generateEventId()) + .author("") + .content( + Content.fromParts( + Part.builder() + .functionCall(FunctionCall.builder().id(callId).name("tool").build()) + .build())) + .build(); + } + + private static ImmutableList ids(List events) { + return events.stream().map(Event::id).collect(toImmutableList()); + } + + private static ImmutableList texts(List events) { + return events.stream().map(Event::stringifyContent).collect(toImmutableList()); + } + + /** + * A plugin that records the events, run errors, and completed runs reported to it by the runner. + */ + private static class RecordingPlugin extends BasePlugin { + final List events = new ArrayList<>(); + final List runErrors = new ArrayList<>(); + int afterRunCount; + + RecordingPlugin() { + super("recording"); + } + + @Override + public Maybe onEventCallback(InvocationContext invocationContext, Event event) { + events.add(event); + return Maybe.empty(); + } + + @Override + public Completable onRunErrorCallback(InvocationContext invocationContext, Throwable error) { + runErrors.add(error); + return Completable.complete(); + } + + @Override + public Completable afterRunCallback(InvocationContext invocationContext) { + afterRunCount++; + return Completable.complete(); + } + } +} diff --git a/core/src/test/java/com/google/adk/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/EdgeTest.java b/core/src/test/java/com/google/adk/workflow/EdgeTest.java index 9fd33b088..2dd0bf85a 100644 --- a/core/src/test/java/com/google/adk/workflow/EdgeTest.java +++ b/core/src/test/java/com/google/adk/workflow/EdgeTest.java @@ -19,6 +19,7 @@ import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertThrows; +import com.google.adk.testing.TestNode; import com.google.common.collect.ImmutableList; import org.junit.Test; import org.junit.runner.RunWith; @@ -27,8 +28,8 @@ @RunWith(JUnit4.class) public final class EdgeTest { - private final Node from = new StubNode("from"); - private final Node to = new StubNode("to"); + private final Node from = TestNode.of("from"); + private final Node to = TestNode.of("to"); @Test public void constructor_withoutRoutes_createsAnUnconditionalEdge() { diff --git a/core/src/test/java/com/google/adk/workflow/EdgesBuilderTest.java b/core/src/test/java/com/google/adk/workflow/EdgesBuilderTest.java index 248dcf4f7..dd1842d4c 100644 --- a/core/src/test/java/com/google/adk/workflow/EdgesBuilderTest.java +++ b/core/src/test/java/com/google/adk/workflow/EdgesBuilderTest.java @@ -19,6 +19,7 @@ import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertThrows; +import com.google.adk.testing.TestNode; import com.google.common.collect.ImmutableList; import java.util.Arrays; import java.util.concurrent.atomic.AtomicReference; @@ -30,10 +31,10 @@ public final class EdgesBuilderTest { private final EdgesBuilder builder = new EdgesBuilder(); - private final Node a = new StubNode("a"); - private final Node b = new StubNode("b"); - private final Node c = new StubNode("c"); - private final Node d = new StubNode("d"); + private final Node a = TestNode.of("a"); + private final Node b = TestNode.of("b"); + private final Node c = TestNode.of("c"); + private final Node d = TestNode.of("d"); private final JoinNode join = JoinNode.builder().name("join").build(); @Test @@ -334,10 +335,10 @@ public void anyOf_everyOverload_dropsRepeatedRoutes() { @Test public void build_documentedExample_buildsItsGraph() { - Node classify = new StubNode("classifier"); - Node category = new StubNode("category"); - Node billing = new StubNode("billing"); - Node summarize = new StubNode("summary"); + Node classify = TestNode.of("classifier"); + Node category = TestNode.of("category"); + Node billing = TestNode.of("billing"); + Node summarize = TestNode.of("summary"); builder .fromStart() diff --git a/core/src/test/java/com/google/adk/workflow/GraphTest.java b/core/src/test/java/com/google/adk/workflow/GraphTest.java index 6b7be3785..7a930c4b6 100644 --- a/core/src/test/java/com/google/adk/workflow/GraphTest.java +++ b/core/src/test/java/com/google/adk/workflow/GraphTest.java @@ -19,6 +19,7 @@ import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertThrows; +import com.google.adk.testing.TestNode; import com.google.common.collect.ImmutableList; import org.junit.Test; import org.junit.runner.RunWith; @@ -32,9 +33,9 @@ public final class GraphTest { @Test public void of_listsEachNodeOnceInTheOrderTheEdgesFirstMentionIt() { - Node a = new StubNode("a"); - Node b = new StubNode("b"); - Node c = new StubNode("c"); + Node a = TestNode.of("a"); + Node b = TestNode.of("b"); + Node c = TestNode.of("c"); ImmutableList edges = ImmutableList.of( new Edge(Start.INSTANCE, a), new Edge(a, c), new Edge(a, b), new Edge(b, c)); @@ -47,21 +48,21 @@ public void of_listsEachNodeOnceInTheOrderTheEdgesFirstMentionIt() { @Test public void terminalNodeNames_areTheNodesWithNoOutgoingEdge() { - Node a = new StubNode("a"); + Node a = TestNode.of("a"); Graph graph = Graph.of( ImmutableList.of( new Edge(Start.INSTANCE, a), - new Edge(a, new StubNode("b")), - new Edge(a, new StubNode("c")))); + new Edge(a, TestNode.of("b")), + new Edge(a, TestNode.of("c")))); assertThat(graph.terminalNodeNames()).containsExactly("b", "c"); } @Test public void node_returnsTheNodeWithThatName() { - Node a = new StubNode("a"); + Node a = TestNode.of("a"); Graph graph = Graph.of(ImmutableList.of(new Edge(Start.INSTANCE, a))); assertThat(graph.node("a")).isSameInstanceAs(a); @@ -69,7 +70,7 @@ public void node_returnsTheNodeWithThatName() { @Test public void node_rejectsAnUnknownName() { - Graph graph = Graph.of(ImmutableList.of(new Edge(Start.INSTANCE, new StubNode("a")))); + Graph graph = Graph.of(ImmutableList.of(new Edge(Start.INSTANCE, TestNode.of("a")))); IllegalArgumentException e = assertThrows(IllegalArgumentException.class, () -> graph.node("missing")); @@ -79,26 +80,26 @@ public void node_rejectsAnUnknownName() { @Test public void nodesTriggeredBy_followsOnlyTheEdgeWithTheEmittedRoute() { - Node router = new StubNode("router"); + Node router = TestNode.of("router"); Graph graph = Graph.of( ImmutableList.of( new Edge(Start.INSTANCE, router), - new Edge(router, new StubNode("yes_node"), YES), - new Edge(router, new StubNode("no_node"), NO))); + new Edge(router, TestNode.of("yes_node"), YES), + new Edge(router, TestNode.of("no_node"), NO))); assertThat(graph.nodesTriggeredBy("router", ImmutableList.of(YES))).containsExactly("yes_node"); } @Test public void nodesTriggeredBy_followsTheEdgeOfEachEmittedRoute() { - Node router = new StubNode("router"); + Node router = TestNode.of("router"); Graph graph = Graph.of( ImmutableList.of( new Edge(Start.INSTANCE, router), - new Edge(router, new StubNode("yes_node"), YES), - new Edge(router, new StubNode("no_node"), NO))); + new Edge(router, TestNode.of("yes_node"), YES), + new Edge(router, TestNode.of("no_node"), NO))); assertThat(graph.nodesTriggeredBy("router", ImmutableList.of(NO, YES))) .containsExactly("yes_node", "no_node") @@ -107,25 +108,25 @@ public void nodesTriggeredBy_followsTheEdgeOfEachEmittedRoute() { @Test public void nodesTriggeredBy_followsAnEdgeOnAnyOfItsRoutes() { - Node router = new StubNode("router"); + Node router = TestNode.of("router"); Graph graph = Graph.of( ImmutableList.of( new Edge(Start.INSTANCE, router), - new Edge(router, new StubNode("either"), ImmutableList.of(YES, NO)))); + new Edge(router, TestNode.of("either"), ImmutableList.of(YES, NO)))); assertThat(graph.nodesTriggeredBy("router", ImmutableList.of(NO))).containsExactly("either"); } @Test public void nodesTriggeredBy_followsTheDefaultEdgeWhenNoRouteMatches() { - Node router = new StubNode("router"); + Node router = TestNode.of("router"); Graph graph = Graph.of( ImmutableList.of( new Edge(Start.INSTANCE, router), - new Edge(router, new StubNode("yes_node"), YES), - new Edge(router, new StubNode("fallback"), Route.Default.INSTANCE))); + new Edge(router, TestNode.of("yes_node"), YES), + new Edge(router, TestNode.of("fallback"), Route.Default.INSTANCE))); assertThat(graph.nodesTriggeredBy("router", ImmutableList.of(new Route.Tag("other")))) .containsExactly("fallback"); @@ -133,13 +134,13 @@ public void nodesTriggeredBy_followsTheDefaultEdgeWhenNoRouteMatches() { @Test public void nodesTriggeredBy_followsTheDefaultEdgeWhenNoRouteIsEmitted() { - Node router = new StubNode("router"); + Node router = TestNode.of("router"); Graph graph = Graph.of( ImmutableList.of( new Edge(Start.INSTANCE, router), - new Edge(router, new StubNode("yes_node"), YES), - new Edge(router, new StubNode("fallback"), Route.Default.INSTANCE))); + new Edge(router, TestNode.of("yes_node"), YES), + new Edge(router, TestNode.of("fallback"), Route.Default.INSTANCE))); assertThat(graph.nodesTriggeredBy("router", ImmutableList.of())).containsExactly("fallback"); } @@ -147,14 +148,14 @@ public void nodesTriggeredBy_followsTheDefaultEdgeWhenNoRouteIsEmitted() { /** Mirrors adk-python's test_get_next_pending_nodes_with_default_route_fan_out. */ @Test public void nodesTriggeredBy_followsEveryDefaultEdgeWhenNoRouteMatches() { - Node router = new StubNode("router"); + Node router = TestNode.of("router"); Graph graph = Graph.of( ImmutableList.of( new Edge(Start.INSTANCE, router), - new Edge(router, new StubNode("named"), YES), - new Edge(router, new StubNode("first"), Route.Default.INSTANCE), - new Edge(router, new StubNode("second"), Route.Default.INSTANCE))); + new Edge(router, TestNode.of("named"), YES), + new Edge(router, TestNode.of("first"), Route.Default.INSTANCE), + new Edge(router, TestNode.of("second"), Route.Default.INSTANCE))); assertThat(graph.nodesTriggeredBy("router", ImmutableList.of(NO))) .containsExactly("first", "second") @@ -163,42 +164,42 @@ public void nodesTriggeredBy_followsEveryDefaultEdgeWhenNoRouteMatches() { @Test public void nodesTriggeredBy_skipsEveryDefaultEdgeWhenARouteMatches() { - Node router = new StubNode("router"); + Node router = TestNode.of("router"); Graph graph = Graph.of( ImmutableList.of( new Edge(Start.INSTANCE, router), - new Edge(router, new StubNode("named"), YES), - new Edge(router, new StubNode("first"), Route.Default.INSTANCE), - new Edge(router, new StubNode("second"), Route.Default.INSTANCE))); + new Edge(router, TestNode.of("named"), YES), + new Edge(router, TestNode.of("first"), Route.Default.INSTANCE), + new Edge(router, TestNode.of("second"), Route.Default.INSTANCE))); assertThat(graph.nodesTriggeredBy("router", ImmutableList.of(YES))).containsExactly("named"); } @Test public void nodesTriggeredBy_followsUnconditionalEdgesAlongsideTheDefaultEdge() { - Node a = new StubNode("a"); + Node a = TestNode.of("a"); Graph graph = Graph.of( ImmutableList.of( new Edge(Start.INSTANCE, a), - new Edge(a, new StubNode("b")), - new Edge(a, new StubNode("c"), YES), - new Edge(a, new StubNode("d"), Route.Default.INSTANCE))); + new Edge(a, TestNode.of("b")), + new Edge(a, TestNode.of("c"), YES), + new Edge(a, TestNode.of("d"), Route.Default.INSTANCE))); assertThat(graph.nodesTriggeredBy("a", ImmutableList.of())).containsExactly("b", "d").inOrder(); } @Test public void nodesTriggeredBy_followsUnconditionalEdgesAlongsideAMatchedRoute() { - Node a = new StubNode("a"); + Node a = TestNode.of("a"); Graph graph = Graph.of( ImmutableList.of( new Edge(Start.INSTANCE, a), - new Edge(a, new StubNode("b")), - new Edge(a, new StubNode("c"), YES), - new Edge(a, new StubNode("d"), Route.Default.INSTANCE))); + new Edge(a, TestNode.of("b")), + new Edge(a, TestNode.of("c"), YES), + new Edge(a, TestNode.of("d"), Route.Default.INSTANCE))); assertThat(graph.nodesTriggeredBy("a", ImmutableList.of(YES))) .containsExactly("b", "c") @@ -207,8 +208,8 @@ public void nodesTriggeredBy_followsUnconditionalEdgesAlongsideAMatchedRoute() { @Test public void nodesTriggeredBy_followsARoutedEdgeBackToAnEarlierNode() { - Node a = new StubNode("a"); - Node b = new StubNode("b"); + Node a = TestNode.of("a"); + Node b = TestNode.of("b"); Graph graph = Graph.of( ImmutableList.of(new Edge(Start.INSTANCE, a), new Edge(a, b), new Edge(b, a, YES))); @@ -218,20 +219,20 @@ public void nodesTriggeredBy_followsARoutedEdgeBackToAnEarlierNode() { @Test public void nodesTriggeredBy_returnsNothingWhenNoRoutedEdgeMatches() { - Node router = new StubNode("router"); + Node router = TestNode.of("router"); Graph graph = Graph.of( ImmutableList.of( - new Edge(Start.INSTANCE, router), new Edge(router, new StubNode("yes_node"), YES))); + new Edge(Start.INSTANCE, router), new Edge(router, TestNode.of("yes_node"), YES))); assertThat(graph.nodesTriggeredBy("router", ImmutableList.of(NO))).isEmpty(); } @Test public void predecessorsOf_returnsTheSourcesOfTheIncomingEdges() { - Node a = new StubNode("a"); - Node b = new StubNode("b"); - Node join = new StubNode("join"); + Node a = TestNode.of("a"); + Node b = TestNode.of("b"); + Node join = TestNode.of("join"); Graph graph = Graph.of( ImmutableList.of( 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..1463c973a --- /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, nodeFailure) -> 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..185019036 --- /dev/null +++ b/core/src/test/java/com/google/adk/workflow/NodeRunnerTest.java @@ -0,0 +1,923 @@ +/* + * 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 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, nodeFailure) -> Completable.fromAction(() -> sentEvents.add(event)); + private final List acknowledgements = new ArrayList<>(); + private final NodeRunner.EventSink acknowledgingSink = + (event, nodeFailure) -> { + sentEvents.add(event); + CompletableSubject acknowledgement = CompletableSubject.create(); + acknowledgements.add(acknowledgement); + return acknowledgement; + }; + private final List eventsInProcessing = new ArrayList<>(); + private final List> processings = new ArrayList<>(); + + @Test + public void run_stampsTheEventWithTheNodeNamePathAndInvocation() { + 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_onASubBranchWithAnOverrideBranch_derivesItFromTheOverrideBranch() { + Context context = + runUnderParent( + new TestNode("greet", (ctx, input) -> Flowable.just("hi")), + /* useSubBranch= */ true, + /* overrideBranch= */ "merged"); + + assertThat(context.branch()).hasValue("merged.greet@1"); + assertThat(getOnlyElement(sentEvents).branch()).hasValue("merged.greet@1"); + } + + @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_asItsParentsOutput_listsBothNodePathsAsItsOutputFor() { + NodeRunner.builder() + .node(new TestNode("worker", (ctx, input) -> Flowable.just("done"))) + .parent(parentContext()) + .runId("2") + .useAsOutput(true) + .build() + .run(/* nodeInput= */ null) + .blockingSubscribe(); + + assertThat(getOnlyElement(sentEvents).nodeInfo().map(NodeInfo::outputFor)) + .hasValue(ImmutableList.of("root@1/worker@2", "root@1")); + } + + @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_partialMessageAsOutputEvent_isNotTheOutput() { + Context context = + runAtRoot( + new TestNode( + "greet", + (ctx, input) -> + Flowable.just( + newEvent() + .partial(true) + .content(MESSAGE) + .nodeInfo(NodeInfo.builder().messageAsOutput(true).build()) + .build()))); + + assertThat(context.requireNodeState().output()).isEmpty(); + } + + @Test + public void run_eventWithoutAnOutput_isNotMarkedAsAnyNodesOutput() { + runAtRoot( + new TestNode("greet", (ctx, input) -> Flowable.just(newEvent().content(MESSAGE).build()))); + + assertThat(getOnlyElement(sentEvents).nodeInfo().map(NodeInfo::outputFor)) + .hasValue(ImmutableList.of()); + } + + @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")); + + Context context = + runAtRoot( + new TestNode( + "router", + (ctx, input) -> { + ctx.requireNodeState().selectRoutes(route); + return Flowable.empty(); + })); + + assertThat(getOnlyElement(sentEvents).actions().route()).hasValue(route); + assertThat(context.requireNodeState().routesEmitted()).isTrue(); + } + + @Test + public void run_outputProducedOnTheNodeState_isSentOnALastEvent() { + Context context = + 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")); + assertThat(context.requireNodeState().outputEmitted()).isTrue(); + } + + @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(); + int sentAfterFirstAcknowledgement = sentEvents.size(); + acknowledgements.get(1).onComplete(); + acknowledgements.get(2).onComplete(); + + assertThat(sentBeforeAcknowledgements).isEqualTo(1); + assertThat(sentAfterFirstAcknowledgement).isEqualTo(2); + 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), this::holdProcessing) + .subscribe(subscriber); + finishProcessing(0); + int emittedBeforeTheCallerReceivesTheFirstEvent = emitted.size(); + subscriber.request(1); + int emittedAfterTheCallerReceivesTheFirstEvent = emitted.size(); + finishProcessing(1); + subscriber.request(1); + + assertThat(emittedBeforeTheCallerReceivesTheFirstEvent).isEqualTo(1); + assertThat(emittedAfterTheCallerReceivesTheFirstEvent).isEqualTo(2); + subscriber.assertValueCount(2).assertComplete(); + } + + @Test + public void runRoot_signalsTheNodeFailureOnlyAfterItsErrorEventIsProcessed() { + IllegalStateException failure = new IllegalStateException("node failed"); + Node node = new TestNode("boom", (ctx, input) -> Flowable.error(failure)); + + TestSubscriber subscriber = + NodeRunner.runRoot(node, invocationContext(/* branch= */ null), this::holdProcessing) + .test(); + subscriber.assertEmpty(); + finishProcessing(0); + + assertThat(getOnlyElement(subscriber.values()).errorCode()).hasValue(ILLEGAL_STATE); + subscriber.assertError(failure); + } + + @Test + public void runRoot_processingTheErrorEventFails_keepsTheNodeFailureAsSuppressed() { + IllegalStateException nodeFailure = new IllegalStateException("node failed"); + IllegalArgumentException processingFailure = new IllegalArgumentException("append failed"); + Node node = new TestNode("boom", (ctx, input) -> Flowable.error(nodeFailure)); + + TestSubscriber subscriber = + NodeRunner.runRoot( + node, + invocationContext(/* branch= */ null), + event -> + event.errorCode().isPresent() + ? Single.error(processingFailure) + : Single.just(event)) + .test(); + + subscriber.assertError(processingFailure); + assertThat(processingFailure.getSuppressed()).asList().containsExactly(nodeFailure); + } + + @Test + public void runRoot_processingAPartialEventFails_keepsTheQueuedNodeFailureAsSuppressed() { + IllegalStateException nodeFailure = new IllegalStateException("node failed"); + IllegalArgumentException processingFailure = new IllegalArgumentException("plugin failed"); + Node node = + new TestNode( + "streamer", + (ctx, input) -> + Flowable.just(newEvent().content(MESSAGE).partial(true).build()) + .concatWith(Flowable.error(nodeFailure))); + + TestSubscriber subscriber = + NodeRunner.runRoot(node, invocationContext(/* branch= */ null), this::holdProcessing) + .test(); + processings.get(0).onError(processingFailure); + + subscriber.assertError(processingFailure); + assertThat(processingFailure.getSuppressed()).asList().containsExactly(nodeFailure); + } + + @Test + public void runRoot_processingAnErrorEventFails_keepsEveryQueuedSiblingFailureAsSuppressed() { + IllegalStateException aFailure = new IllegalStateException("a failed"); + IllegalStateException bFailure = new IllegalStateException("b failed"); + IllegalArgumentException processingFailure = new IllegalArgumentException("append failed"); + Workflow workflow = + Workflow.builder() + .name("wf") + .edges( + new Edge( + Start.INSTANCE, new TestNode("a", (ctx, input) -> Flowable.error(aFailure))), + new Edge( + Start.INSTANCE, new TestNode("b", (ctx, input) -> Flowable.error(bFailure)))) + .build(); + + TestSubscriber subscriber = + NodeRunner.runRoot(workflow, invocationContext(/* branch= */ null), this::holdProcessing) + .test(); + processings.get(0).onError(processingFailure); + + subscriber.assertError(processingFailure); + assertThat(processingFailure.getSuppressed()).asList().containsExactly(aFailure, bFailure); + } + + @Test + public void runRoot_processingANonErrorEventFails_keepsAQueuedSiblingFailureAsSuppressed() { + IllegalStateException bFailure = new IllegalStateException("b failed"); + IllegalArgumentException processingFailure = new IllegalArgumentException("append failed"); + Workflow workflow = + Workflow.builder() + .name("wf") + .edges( + new Edge( + Start.INSTANCE, + new TestNode( + "a", (ctx, input) -> Flowable.just(newEvent().content(MESSAGE).build()))), + new Edge( + Start.INSTANCE, new TestNode("b", (ctx, input) -> Flowable.error(bFailure)))) + .build(); + + TestSubscriber subscriber = + NodeRunner.runRoot(workflow, invocationContext(/* branch= */ null), this::holdProcessing) + .test(); + Event inProcessing = getOnlyElement(eventsInProcessing); + processings.get(0).onError(processingFailure); + + assertThat(inProcessing.errorCode()).isEmpty(); + subscriber.assertError(processingFailure); + assertThat(processingFailure.getSuppressed()).asList().containsExactly(bFailure); + } + + @Test + public void runRoot_emitsTheEventThatProcessingReturns() { + Node node = + new TestNode("solo", (ctx, input) -> Flowable.just(newEvent().content(MESSAGE).build())); + Event replacement = newEvent().author("plugin").content(MESSAGE).build(); + + TestSubscriber subscriber = + NodeRunner.runRoot( + node, invocationContext(/* branch= */ null), event -> Single.just(replacement)) + .test(); + + subscriber.assertComplete(); + assertThat(getOnlyElement(subscriber.values())).isSameInstanceAs(replacement); + } + + @Test + public void runRoot_errorFromTheNode_failsTheRunWithoutAnErrorEvent() { + AssertionError fatal = new AssertionError("fatal"); + Node node = new TestNode("fatal", (ctx, input) -> Flowable.error(fatal)); + + TestSubscriber subscriber = + NodeRunner.runRoot(node, invocationContext(/* branch= */ null), this::holdProcessing) + .test(); + + subscriber.assertError(fatal); + assertThat(eventsInProcessing).isEmpty(); + } + + @Test + public void runRoot_partialEvent_doesNotHoldBackTheNode() { + List emitted = new ArrayList<>(); + Node node = + new TestNode( + "streamer", + (ctx, input) -> + Flowable.just( + newEvent().content(MESSAGE).partial(true).build(), + newEvent().content(MESSAGE).build()) + .doOnNext(emitted::add)); + + TestSubscriber subscriber = + NodeRunner.runRoot(node, invocationContext(/* branch= */ null), this::holdProcessing) + .test(); + int emittedWhileThePartialEventIsProcessed = emitted.size(); + finishProcessing(0); + finishProcessing(1); + + assertThat(emittedWhileThePartialEventIsProcessed).isEqualTo(2); + subscriber.assertValueCount(2).assertComplete(); + } + + @CanIgnoreReturnValue + private Context runAtRoot(Node node) { + return rootRunner(node, recordingSink).run(/* nodeInput= */ null).blockingGet(); + } + + private static NodeRunner rootRunner(Node node, NodeRunner.EventSink eventSink) { + return NodeRunner.builder() + .node(node) + .invocationContext(invocationContext(/* branch= */ null)) + .eventSink(eventSink) + .build(); + } + + private Single holdProcessing(Event event) { + eventsInProcessing.add(event); + SingleSubject processing = SingleSubject.create(); + processings.add(processing); + return processing; + } + + private void finishProcessing(int index) { + processings.get(index).onSuccess(eventsInProcessing.get(index)); + } + + private Context runUnderParent(Node node, boolean useSubBranch, @Nullable String overrideBranch) { + return NodeRunner.builder() + .node(node) + .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(); + } +} diff --git a/core/src/test/java/com/google/adk/workflow/SchedulerTest.java b/core/src/test/java/com/google/adk/workflow/SchedulerTest.java new file mode 100644 index 000000000..71a14fa7b --- /dev/null +++ b/core/src/test/java/com/google/adk/workflow/SchedulerTest.java @@ -0,0 +1,651 @@ +/* + * 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.Iterables.getLast; +import static com.google.common.collect.MoreCollectors.onlyElement; +import static com.google.common.truth.Truth.assertThat; + +import com.google.adk.agents.Context; +import com.google.adk.agents.InvocationContext; +import com.google.adk.events.Event; +import com.google.adk.testing.TestNode; +import com.google.adk.testing.TestUtils; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; +import com.google.errorprone.annotations.CanIgnoreReturnValue; +import com.google.genai.types.FinishReason; +import io.reactivex.rxjava3.core.Completable; +import io.reactivex.rxjava3.core.Flowable; +import io.reactivex.rxjava3.observers.TestObserver; +import io.reactivex.rxjava3.subjects.CompletableSubject; +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.jspecify.annotations.Nullable; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public final class SchedulerTest { + + private static final Start START = Start.INSTANCE; + private static final Route YES = new Route.Tag("yes"); + private static final Route NO = new Route.Tag("no"); + + private final List sentEvents = new ArrayList<>(); + private final NodeRunner.EventSink recordingSink = + (event, nodeFailure) -> Completable.fromAction(() -> sentEvents.add(event)); + private final AtomicInteger activeRuns = new AtomicInteger(); + private final AtomicInteger peakRuns = new AtomicInteger(); + + @Test + public void aSingleSuccessor_keepsItsPredecessorsBranch() { + Node a = emitter("a", "A"); + Node b = emitter("b", "B"); + + run(workflow(new Edge(START, a), new Edge(a, b))); + + assertThat(eventAt("wf@1/a@1").branch()).hasValue("root"); + assertThat(eventAt("wf@1/b@1").branch()).hasValue("root"); + } + + @Test + public void aJoin_receivesEveryPredecessorOutputKeyedByNodeName() { + Node a = emitter("a", "A"); + Node b = emitter("b", "B"); + Node join = JoinNode.builder().name("join").build(); + + run(workflow(new Edge(START, a), new Edge(START, b), new Edge(a, join), new Edge(b, join))); + + assertThat(eventAt("wf@1/join@1").output()).hasValue(ImmutableMap.of("a", "A", "b", "B")); + } + + @Test + public void aJoinWithOnePredecessor_keepsThatPredecessorsBranch() { + Node a = emitter("a", "A"); + Node join = JoinNode.builder().name("join").build(); + + run(workflow(new Edge(START, a), new Edge(a, join))); + + assertThat(eventAt("wf@1/join@1").output()).hasValue(ImmutableMap.of("a", "A")); + assertThat(eventAt("wf@1/join@1").branch()).hasValue("root"); + } + + @Test + public void aJoin_reMergesItsForkedPredecessorsToTheirCommonBranch() { + Node a = emitter("a", "A"); + Node b = emitter("b", "B"); + Node join = JoinNode.builder().name("join").build(); + + run(workflow(new Edge(START, a), new Edge(START, b), new Edge(a, join), new Edge(b, join))); + + assertThat(eventAt("wf@1/a@1").branch()).hasValue("root.a@1"); + assertThat(eventAt("wf@1/b@1").branch()).hasValue("root.b@1"); + assertThat(eventAt("wf@1/join@1").branch()).hasValue("root"); + } + + @Test + public void aJoin_reMergesToTheBranchItsPredecessorsForkedFromBelowTheRoot() { + Node fork = emitter("fork", "F"); + Node a = emitter("a", "A"); + Node b = emitter("b", "B"); + Node join = JoinNode.builder().name("join").build(); + + run( + workflow( + new Edge(START, fork), + new Edge(START, TestNode.of("idle")), + new Edge(fork, a), + new Edge(fork, b), + new Edge(a, join), + new Edge(b, join))); + + assertThat(eventAt("wf@1/a@1").branch()).hasValue("root.fork@1.a@1"); + assertThat(eventAt("wf@1/b@1").branch()).hasValue("root.fork@1.b@1"); + assertThat(eventAt("wf@1/join@1").branch()).hasValue("root.fork@1"); + } + + @Test + public void aJoinFedByTheStartNode_receivesTheWorkflowInputBesideItsOtherPredecessors() { + Node x = emitter("x", "X"); + Node join = JoinNode.builder().name("join").build(); + + runner(workflow(new Edge(START, x), new Edge(START, join), new Edge(x, join)), "root") + .run("IN") + .blockingSubscribe(); + + assertThat(eventAt("wf@1/join@1").output()) + .hasValue(ImmutableMap.of(Start.NAME, "IN", "x", "X")); + assertThat(eventAt("wf@1/join@1").branch()).hasValue("root"); + } + + @Test + public void aJoinWithoutARootBranch_reMergesToNoBranch() { + Node a = emitter("a", "A"); + Node b = emitter("b", "B"); + Node join = JoinNode.builder().name("join").build(); + + run( + workflow(new Edge(START, a), new Edge(START, b), new Edge(a, join), new Edge(b, join)), + /* branch= */ null); + + assertThat(eventAt("wf@1/a@1").branch()).hasValue("a@1"); + assertThat(eventAt("wf@1/join@1").branch()).isEmpty(); + } + + @Test + public void theSoleTerminalNodesOutput_becomesTheWorkflowsOutputWithoutAnotherEvent() { + Node a = emitter("a", "A"); + Node b = emitter("b", "B"); + + Context context = run(workflow(new Edge(START, a), new Edge(a, b))); + + assertThat(context.requireNodeState().output()).hasValue("B"); + assertThat(eventAt("wf@1/b@1").nodeInfo().map(NodeInfo::outputFor)) + .hasValue(ImmutableList.of("wf@1/b@1", "wf@1")); + assertThat(sentEvents).hasSize(2); + } + + @Test + public void twoTerminalNodesProducingOutput_failTheWorkflow() { + Workflow workflow = + workflow(new Edge(START, emitter("a", "A")), new Edge(START, emitter("b", "B"))); + + Context context = runner(workflow, "root").run(/* nodeInput= */ null).blockingGet(); + + assertThat(context.requireNodeState().failure().orElseThrow()) + .isInstanceOf(WorkflowConfigurationException.class); + Event errorEvent = getLast(sentEvents); + assertThat(errorEvent.nodeInfo().map(NodeInfo::path)).hasValue("wf@1"); + assertThat(errorEvent.errorCode()).hasValue(new FinishReason("WorkflowConfigurationException")); + } + + @Test + public void theEventsOfTheNodesInside_areAttributedToTheWorkflow() { + run(workflow(new Edge(START, emitter("a", "A")))); + + assertThat(eventAt("wf@1/a@1").author()).isEqualTo("wf"); + } + + @Test + public void anAgentNode_keepsItsAuthorAndFeedsItsSuccessor() { + Node agent = + TestUtils.createSubAgent( + "worker", Event.builder().id("e1").author("worker").output("FROM_AGENT").build()); + List received = new ArrayList<>(); + + run(workflow(new Edge(START, agent), new Edge(agent, recorder("next", received)))); + + assertThat(eventAt("wf@1/worker@1").author()).isEqualTo("worker"); + assertThat(received).containsExactly("FROM_AGENT"); + } + + @Test + public void aNestedWorkflow_handsItsTerminalOutputToTheOuterSuccessor() { + Workflow inner = + Workflow.builder().name("inner").edges(new Edge(START, emitter("leaf", "INNER"))).build(); + List received = new ArrayList<>(); + + run(workflow(new Edge(START, inner), new Edge(inner, recorder("next", received)))); + + assertThat(received).containsExactly("INNER"); + assertThat(eventAt("wf@1/inner@1/leaf@1").nodeInfo().map(NodeInfo::outputFor)) + .hasValue(ImmutableList.of("wf@1/inner@1/leaf@1", "wf@1/inner@1")); + } + + @Test + public void anEmittedRoute_triggersOnlyTheMatchingSuccessor() { + Node router = router("router", YES); + + run( + workflow( + new Edge(START, router), + new Edge(router, emitter("taken", "T"), YES), + new Edge(router, emitter("skipped", "S"), NO))); + + assertThat(eventsAt("wf@1/taken@1")).hasSize(1); + assertThat(eventsAt("wf@1/skipped@1")).isEmpty(); + } + + @Test + public void anUnmatchedRoute_triggersEveryDefaultSuccessor() { + Node router = router("router", new Route.Tag("unmatched")); + Node b = emitter("b", "B"); + Node c = emitter("c", "C"); + Node join = JoinNode.builder().name("join").build(); + + run( + workflow( + new Edge(START, router), + new Edge(router, b, Route.Default.INSTANCE), + new Edge(router, c, Route.Default.INSTANCE), + new Edge(b, join), + new Edge(c, join))); + + assertThat(eventAt("wf@1/join@1").output()).hasValue(ImmutableMap.of("b", "B", "c", "C")); + } + + @Test + public void maxConcurrency_capsHowManyNodesRunAtOnce() { + CompletableSubject gate = CompletableSubject.create(); + Workflow workflow = + Workflow.builder() + .name("wf") + .maxConcurrency(1) + .edges( + new Edge(START, gated("a", gate)), + new Edge(START, gated("b", gate)), + new Edge(START, gated("c", gate))) + .build(); + + TestObserver observer = runner(workflow, "root").run(/* nodeInput= */ null).test(); + gate.onComplete(); + + observer.assertComplete(); + assertThat(peakRuns.get()).isEqualTo(1); + } + + @Test + public void withoutAMaxConcurrency_forkedNodesRunTogether() { + CompletableSubject gate = CompletableSubject.create(); + Workflow workflow = + workflow( + new Edge(START, gated("a", gate)), + new Edge(START, gated("b", gate)), + new Edge(START, gated("c", gate))); + + TestObserver observer = runner(workflow, "root").run(/* nodeInput= */ null).test(); + gate.onComplete(); + + observer.assertComplete(); + assertThat(peakRuns.get()).isEqualTo(3); + } + + @Test + public void aRunningNode_takesItsNextTriggerOnlyAfterItFinishes() { + CompletableSubject gate = CompletableSubject.create(); + CompletableSubject slowGate = CompletableSubject.create(); + AtomicInteger runs = new AtomicInteger(); + Node a = emitter("a", "A"); + Node g = new TestNode("g", (ctx, input) -> gate.andThen(Flowable.just("G"))); + Node slow = + new TestNode( + "slow", + (ctx, input) -> + Flowable.defer( + () -> { + runs.incrementAndGet(); + peakRuns.accumulateAndGet(activeRuns.incrementAndGet(), Math::max); + return slowGate.doOnComplete(activeRuns::decrementAndGet).toFlowable(); + })); + + TestObserver observer = + runner( + workflow( + new Edge(START, a), new Edge(START, g), new Edge(a, slow), new Edge(g, slow)), + "root") + .run(/* nodeInput= */ null) + .test(); + gate.onComplete(); + slowGate.onComplete(); + + observer.assertComplete(); + assertThat(runs.get()).isEqualTo(2); + assertThat(peakRuns.get()).isEqualTo(1); + } + + @Test + public void aNodeWaitingOnAnswers_keepsItsLaterTriggersQueued() { + AtomicInteger runs = new AtomicInteger(); + Node asker = + new TestNode( + "asker", + (ctx, input) -> { + int run = runs.incrementAndGet(); + return Flowable.just( + Event.builder() + .id("e" + run) + .author("") + .longRunningToolIds(ImmutableSet.of("call-" + run)) + .build()); + }); + Node a = emitter("a", "A"); + Node b = emitter("b", "B"); + + Context context = + run( + Workflow.builder() + .name("wf") + .maxConcurrency(1) + .edges( + new Edge(START, a), new Edge(START, b), new Edge(a, asker), new Edge(b, asker)) + .build()); + + assertThat(runs.get()).isEqualTo(1); + assertThat(context.requireNodeState().interruptIds()).containsExactly("call-1"); + } + + @Test + public void aJoinUnderAConcurrencyCap_runsOnceWithEveryPredecessorOutput() { + Node a = emitter("a", "A"); + Node b = emitter("b", "B"); + Node c = emitter("c", "C"); + Node join = JoinNode.builder().name("join").build(); + + run( + Workflow.builder() + .name("wf") + .maxConcurrency(1) + .edges( + new Edge(START, a), + new Edge(START, b), + new Edge(START, c), + new Edge(a, join), + new Edge(b, join), + new Edge(c, join)) + .build()); + + assertThat(eventAt("wf@1/join@1").output()) + .hasValue(ImmutableMap.of("a", "A", "b", "B", "c", "C")); + } + + @Test + public void aWaitForOutputNode_waitsUntilATriggerProducesItsOutput() { + Node a = emitter("a", "A"); + Node b = emitter("b", "B"); + LateProducer late = new LateProducer(/* readyOn= */ 2); + List received = new ArrayList<>(); + + run( + workflow( + new Edge(START, a), + new Edge(START, b), + new Edge(a, late), + new Edge(b, late), + new Edge(late, recorder("next", received)))); + + assertThat(late.runs).isEqualTo(2); + assertThat(received).containsExactly("ready"); + } + + @Test + public void aLoopThatRetriggersAJoinsPredecessors_givesTheJoinOneSetOfOutputsPerIteration() { + Node a = new CountingEmitter("a"); + Node b = new CountingEmitter("b"); + LoopingJoin join = new LoopingJoin(); + + run( + Workflow.builder() + .name("wf") + .maxConcurrency(1) + .edges( + new Edge(START, a), + new Edge(START, b), + new Edge(a, join), + new Edge(b, join), + new Edge(join, a, YES), + new Edge(join, b, YES)) + .build()); + + assertThat(join.seen) + .containsExactly( + ImmutableMap.of("a", "a1", "b", "b1"), ImmutableMap.of("a", "a2", "b", "b2")) + .inOrder(); + assertThat(eventsAt("wf@1/a@2")).hasSize(1); + } + + @Test + public void aFailedNode_failsTheWorkflowAndCancelsTheNodesStillRunning() { + IllegalStateException failure = new IllegalStateException("node failed"); + AtomicBoolean canceled = new AtomicBoolean(); + Node slow = + new TestNode( + "slow", + (ctx, input) -> + CompletableSubject.create().toFlowable().doOnCancel(() -> canceled.set(true))); + Node boom = new TestNode("boom", (ctx, input) -> Flowable.error(failure)); + List received = new ArrayList<>(); + + TestObserver observer = + runner( + workflow( + new Edge(START, slow), + new Edge(START, boom), + new Edge(boom, recorder("next", received))), + "root") + .run(/* nodeInput= */ null) + .test(); + + observer.assertComplete(); + assertThat(observer.values().get(0).requireNodeState().failure()).hasValue(failure); + assertThat(eventAt("wf@1/boom@1").errorCode()) + .hasValue(new FinishReason("IllegalStateException")); + assertThat(canceled.get()).isTrue(); + assertThat(received).isEmpty(); + } + + @Test + public void aNodeFinishingAfterAFailure_doesNotStartItsSuccessor() { + CompletableSubject gateB = CompletableSubject.create(); + CompletableSubject gateC = CompletableSubject.create(); + // Canceling either waiting node finishes the other, so one finishes during failure handling. + Node b = + new TestNode( + "b", (ctx, input) -> gateB.andThen(Flowable.just("B")).doOnCancel(gateC::onComplete)); + Node c = + new TestNode( + "c", (ctx, input) -> gateC.andThen(Flowable.just("C")).doOnCancel(gateB::onComplete)); + IllegalStateException failure = new IllegalStateException("boom"); + Node failer = new TestNode("failer", (ctx, input) -> Flowable.error(failure)); + List received = new ArrayList<>(); + Workflow workflow = + workflow( + new Edge(START, failer), + new Edge(START, b), + new Edge(START, c), + new Edge(b, recorder("afterB", received)), + new Edge(c, recorder("afterC", received))); + + Context context = runner(workflow, "root").run(/* nodeInput= */ null).blockingGet(); + + assertThat(context.requireNodeState().failure()).hasValue(failure); + assertThat(received).isEmpty(); + } + + @Test + public void anErrorFromANode_failsTheWorkflowRun() { + AssertionError fatal = new AssertionError("fatal"); + + runner( + workflow(new Edge(START, new TestNode("fatal", (ctx, input) -> Flowable.error(fatal)))), + "root") + .run(/* nodeInput= */ null) + .test() + .assertError(fatal); + } + + @Test + public void aNodeWaitingOnAnswers_endsItsBranchAndSurfacesItsInterruptsOnTheWorkflow() { + Node asker = + new TestNode( + "asker", + (ctx, input) -> + Flowable.just( + Event.builder() + .id("e1") + .author("") + .longRunningToolIds(ImmutableSet.of("call-1")) + .build())); + List received = new ArrayList<>(); + + Context context = + run(workflow(new Edge(START, asker), new Edge(asker, recorder("next", received)))); + + assertThat(context.requireNodeState().interruptIds()).containsExactly("call-1"); + assertThat(received).isEmpty(); + } + + @CanIgnoreReturnValue + private Context run(Workflow workflow) { + return run(workflow, "root"); + } + + /** Runs {@code workflow} at the root on {@code branch} and asserts that it did not fail. */ + @CanIgnoreReturnValue + private Context run(Workflow workflow, @Nullable String branch) { + Context context = runner(workflow, branch).run(/* nodeInput= */ null).blockingGet(); + assertThat(context.requireNodeState().failure()).isEmpty(); + return context; + } + + /** Returns a runner of {@code workflow} at the root of a run on {@code branch}. */ + private NodeRunner runner(Workflow workflow, @Nullable String branch) { + InvocationContext invocationContext = + TestUtils.createInvocationContext(TestUtils.createRootAgent()).toBuilder() + .branch(branch) + .build(); + return NodeRunner.builder() + .node(workflow) + .invocationContext(invocationContext) + .eventSink(recordingSink) + .build(); + } + + private Event eventAt(String nodePath) { + return eventsAt(nodePath).stream().collect(onlyElement()); + } + + private ImmutableList eventsAt(String nodePath) { + return sentEvents.stream() + .filter(event -> event.nodeInfo().map(NodeInfo::path).equals(Optional.of(nodePath))) + .collect(ImmutableList.toImmutableList()); + } + + /** Returns a node that stays running until {@code gate} completes, counting overlapping runs. */ + private Node gated(String name, CompletableSubject gate) { + return new TestNode( + name, + (ctx, input) -> + Flowable.defer( + () -> { + peakRuns.accumulateAndGet(activeRuns.incrementAndGet(), Math::max); + // Leaves before completing, since the scheduler may start the next node then. + return gate.toFlowable().doOnComplete(activeRuns::decrementAndGet); + })); + } + + private static Workflow workflow(Edge... edges) { + return Workflow.builder().name("wf").edges(edges).build(); + } + + private static Node emitter(String name, Object output) { + return new TestNode(name, (ctx, input) -> Flowable.just(output)); + } + + private static Node recorder(String name, List received) { + return new TestNode( + name, + (ctx, input) -> { + received.add(input); + return Flowable.empty(); + }); + } + + private static Node router(String name, Route route) { + return new TestNode( + name, + (ctx, input) -> { + ctx.requireNodeState().selectRoutes(ImmutableList.of(route)); + return Flowable.empty(); + }); + } + + /** A node that waits for an output and produces one only from its {@code readyOn}-th run. */ + private static final class LateProducer implements Node { + private final int readyOn; + private int runs; + + LateProducer(int readyOn) { + this.readyOn = readyOn; + } + + @Override + public String name() { + return "late"; + } + + @Override + public boolean waitForOutput() { + return true; + } + + @Override + public Flowable runNode(Context context, @Nullable Object nodeInput) { + return ++runs >= readyOn ? Flowable.just("ready") : Flowable.empty(); + } + } + + /** A node whose output names its run, so the iterations of a loop are distinguishable. */ + private static final class CountingEmitter implements Node { + private final String name; + private int runs; + + CountingEmitter(String name) { + this.name = name; + } + + @Override + public String name() { + return name; + } + + @Override + public Flowable runNode(Context context, @Nullable Object nodeInput) { + return Flowable.just(name + ++runs); + } + } + + /** A join that routes back to its predecessors once, recording every input it receives. */ + private static final class LoopingJoin implements Node { + private final List<@Nullable Object> seen = new ArrayList<>(); + + @Override + public String name() { + return "join"; + } + + @Override + public boolean requiresAllPredecessors() { + return true; + } + + @Override + public Flowable runNode(Context context, @Nullable Object nodeInput) { + seen.add(nodeInput); + if (seen.size() == 1) { + context.requireNodeState().selectRoutes(ImmutableList.of(YES)); + return Flowable.empty(); + } + return Flowable.just(nodeInput); + } + } +} diff --git a/core/src/test/java/com/google/adk/workflow/WorkflowTest.java b/core/src/test/java/com/google/adk/workflow/WorkflowTest.java index cf1f90aa2..e4ea1318b 100644 --- a/core/src/test/java/com/google/adk/workflow/WorkflowTest.java +++ b/core/src/test/java/com/google/adk/workflow/WorkflowTest.java @@ -20,6 +20,7 @@ import static org.junit.Assert.assertThrows; import com.google.adk.agents.CallbackContext; +import com.google.adk.testing.TestNode; import com.google.adk.testing.TestUtils; import com.google.common.collect.ImmutableList; import com.google.genai.types.Schema; @@ -35,8 +36,8 @@ public final class WorkflowTest { @Test public void build_buildsTheGraphFromTheEdges() { - Node a = new StubNode("a"); - Node b = new StubNode("b"); + Node a = TestNode.of("a"); + Node b = TestNode.of("b"); ImmutableList edges = ImmutableList.of(new Edge(Start.INSTANCE, a), new Edge(a, b)); Workflow workflow = Workflow.builder().name("flow").edges(edges).build(); @@ -103,9 +104,9 @@ public void build_setsEveryProperty() { @Test public void edges_withVarargs_setsTheEdgesInOrder() { - Node a = new StubNode("a"); + Node a = TestNode.of("a"); Edge first = new Edge(Start.INSTANCE, a); - Edge second = new Edge(a, new StubNode("b")); + Edge second = new Edge(a, TestNode.of("b")); Workflow workflow = Workflow.builder().name("flow").edges(first, second).build(); @@ -114,8 +115,8 @@ public void edges_withVarargs_setsTheEdgesInOrder() { @Test public void edges_withAConsumer_setsTheEdgesItAdds() { - Node a = new StubNode("a"); - Node b = new StubNode("b"); + Node a = TestNode.of("a"); + Node b = TestNode.of("b"); Workflow workflow = Workflow.builder().name("flow").edges(g -> g.fromStart().then(a).then(b)).build(); @@ -189,7 +190,7 @@ public void build_rejectsAMissingName() { @Test public void aWorkflowCanBeANodeOfAnotherWorkflow() { Workflow inner = - Workflow.builder().name("inner").edges(new Edge(Start.INSTANCE, new StubNode("a"))).build(); + Workflow.builder().name("inner").edges(new Edge(Start.INSTANCE, TestNode.of("a"))).build(); Workflow outer = Workflow.builder().name("outer").edges(new Edge(Start.INSTANCE, inner)).build(); @@ -198,17 +199,13 @@ public void aWorkflowCanBeANodeOfAnotherWorkflow() { } @Test - public void runNode_failsAsNotSupportedYet() { - Workflow workflow = - Workflow.builder().name("flow").edges(new Edge(Start.INSTANCE, new StubNode("a"))).build(); + public void runNode_withoutEdges_completesWithoutEmitting() { + Workflow workflow = Workflow.builder().name("flow").build(); CallbackContext context = new CallbackContext( TestUtils.createInvocationContext(TestUtils.createRootAgent()), /* eventActions= */ null); - workflow - .runNode(context, /* nodeInput= */ null) - .test() - .assertError(UnsupportedOperationException.class); + workflow.runNode(context, /* nodeInput= */ null).test().assertComplete().assertNoValues(); } }