From cecad8f49e7ce58aea790ab598753f1475b685dd Mon Sep 17 00:00:00 2001 From: Damian Momot Date: Wed, 21 Oct 2026 03:06:17 -0700 Subject: [PATCH] feat: run a workflow graph PiperOrigin-RevId: 1003793324 --- .../java/com/google/adk/runner/Runner.java | 35 + .../adk/workflow/NodeExecutionState.java | 29 + .../com/google/adk/workflow/NodeRunner.java | 18 +- .../com/google/adk/workflow/NodeState.java | 35 + .../com/google/adk/workflow/NodeStatus.java | 32 + .../com/google/adk/workflow/Scheduler.java | 336 +++++++++ .../java/com/google/adk/workflow/Trigger.java | 35 + .../com/google/adk/workflow/TriggerQueue.java | 58 ++ .../com/google/adk/workflow/Workflow.java | 20 +- .../WorkflowConfigurationException.java | 28 + .../google/adk/runner/RunnerRootNodeTest.java | 115 +++ .../adk/workflow/NodeExecutionStateTest.java | 14 + .../google/adk/workflow/NodeRunnerTest.java | 66 ++ .../google/adk/workflow/SchedulerTest.java | 676 ++++++++++++++++++ .../com/google/adk/workflow/WorkflowTest.java | 10 +- 15 files changed, 1495 insertions(+), 12 deletions(-) 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 create mode 100644 core/src/main/java/com/google/adk/workflow/WorkflowConfigurationException.java create mode 100644 core/src/test/java/com/google/adk/workflow/SchedulerTest.java 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 64f84cf32..96ae41f96 100644 --- a/core/src/main/java/com/google/adk/runner/Runner.java +++ b/core/src/main/java/com/google/adk/runner/Runner.java @@ -56,8 +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; @@ -81,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; @@ -474,10 +477,42 @@ public PluginManager pluginManager() { public Completable close() { List completables = new ArrayList<>(); completables.add(agent.close()); + if (node != null) { + Set nodeAgents = new LinkedHashSet<>(); + collectAgents(node, nodeAgents); + // Closing an agent closes its sub-agents, so skip agents whose ancestor is closed here. + nodeAgents.stream() + .filter(nodeAgent -> !hasAncestorIn(nodeAgent, 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); + } + } + } + + /** Returns whether {@code agents} contains an ancestor of {@code agent}. */ + private static boolean hasAncestorIn(BaseAgent agent, Set agents) { + for (BaseAgent ancestor = agent.parentAgent(); + ancestor != null; + ancestor = ancestor.parentAgent()) { + if (agents.contains(ancestor)) { + return true; + } + } + return false; + } + /** * Prepares and creates the user-authored {@link Event}, saving any input blobs as artifacts and * attaching state deltas if configured. diff --git a/core/src/main/java/com/google/adk/workflow/NodeExecutionState.java b/core/src/main/java/com/google/adk/workflow/NodeExecutionState.java index 1801e52f2..3c191ff67 100644 --- a/core/src/main/java/com/google/adk/workflow/NodeExecutionState.java +++ b/core/src/main/java/com/google/adk/workflow/NodeExecutionState.java @@ -25,8 +25,10 @@ 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; @@ -43,6 +45,8 @@ public final class NodeExecutionState { 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; @@ -66,6 +70,10 @@ private NodeExecutionState(Builder builder) { 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; } @@ -156,6 +164,19 @@ 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; @@ -194,6 +215,7 @@ static final class Builder { private NodeRunner.@Nullable EventSink eventSink; private @Nullable Context parent; private @Nullable String runId; + private boolean useAsOutput; private Builder() {} @@ -228,6 +250,13 @@ Builder runId(String 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. * diff --git a/core/src/main/java/com/google/adk/workflow/NodeRunner.java b/core/src/main/java/com/google/adk/workflow/NodeRunner.java index 4c7a0a31b..e2c6ae389 100644 --- a/core/src/main/java/com/google/adk/workflow/NodeRunner.java +++ b/core/src/main/java/com/google/adk/workflow/NodeRunner.java @@ -26,7 +26,6 @@ import com.google.adk.annotations.Experimental; import com.google.adk.events.Event; import com.google.adk.events.EventActions; -import com.google.common.collect.ImmutableList; import com.google.errorprone.annotations.CanIgnoreReturnValue; import com.google.genai.types.FinishReason; import io.reactivex.rxjava3.core.Completable; @@ -68,6 +67,7 @@ interface EventSink { 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; @@ -89,6 +89,7 @@ private NodeRunner(Builder builder) { this.node = builder.node; this.parent = builder.parent; this.runId = builder.runId; + this.useAsOutput = builder.useAsOutput; this.useSubBranch = builder.useSubBranch; this.overrideBranch = builder.overrideBranch; } @@ -185,6 +186,7 @@ Single run(@Nullable Object nodeInput) { .eventSink(eventSink) .parent(parent) .runId(runId) + .useAsOutput(useAsOutput) .build()); return new Activation(context).run(nodeInput).toSingleDefault(context); }); @@ -355,8 +357,7 @@ private Event stamp(Event event) { .nodeInfo( nodeInfo.toBuilder() .path(nodeState.nodePath()) - .outputFor( - carriesOutput ? ImmutableList.of(nodeState.nodePath()) : nodeInfo.outputFor()) + .outputFor(carriesOutput ? nodeState.outputFor() : nodeInfo.outputFor()) .build()) .build(); } @@ -419,6 +420,7 @@ static final class Builder { private @Nullable EventSink eventSink; private @Nullable Context parent; private String runId = "1"; + private boolean useAsOutput; private boolean useSubBranch; private @Nullable String overrideBranch; @@ -459,6 +461,16 @@ Builder runId(String 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. 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..cc335c906 --- /dev/null +++ b/core/src/main/java/com/google/adk/workflow/Scheduler.java @@ -0,0 +1,336 @@ +/* + * 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 workflowState; + 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.workflowState = 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()) { + workflowState.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 progress = nodeStates.get(nodeName); + if (progress != null + && progress.status == NodeStatus.WAITING + && !progress.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 = workflowState.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 progress = nodeStates.get(nodeName); + NodeExecutionState childState = childContext.requireNodeState(); + ImmutableSet childInterruptIds = childState.interruptIds(); + if (!childInterruptIds.isEmpty()) { + progress.status = NodeStatus.WAITING; + progress.interrupts = childInterruptIds.asList(); + interruptIds.addAll(childInterruptIds); + return; + } + if (graph.node(nodeName).waitForOutput() + && childState.output().isEmpty() + && childState.selectedRoutes().isEmpty()) { + progress.status = NodeStatus.WAITING; + return; + } + progress.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 progress = nodeStates.get(predecessor); + boolean done = + predecessor.equals(Start.NAME) + || (progress != null + && progress.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()) { + workflowState.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) { + workflowState.produceOutput(nodeOutputs.get(terminalOutputs.get(0))); + // The terminal node's event already carried the output for the workflow as well. + workflowState.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..a14a8382c 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,25 @@ 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. If more than one + * terminal node produces an output, the returned {@code Flowable} fails with {@link + * WorkflowConfigurationException}. + */ @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/main/java/com/google/adk/workflow/WorkflowConfigurationException.java b/core/src/main/java/com/google/adk/workflow/WorkflowConfigurationException.java new file mode 100644 index 000000000..b75e21b20 --- /dev/null +++ b/core/src/main/java/com/google/adk/workflow/WorkflowConfigurationException.java @@ -0,0 +1,28 @@ +/* + * 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; + +/** 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/runner/RunnerRootNodeTest.java b/core/src/test/java/com/google/adk/runner/RunnerRootNodeTest.java index a931e13ef..801f251b3 100644 --- a/core/src/test/java/com/google/adk/runner/RunnerRootNodeTest.java +++ b/core/src/test/java/com/google/adk/runner/RunnerRootNodeTest.java @@ -16,17 +16,22 @@ package com.google.adk.runner; +import static com.google.adk.testing.TestUtils.createTestAgentBuilder; +import static com.google.adk.testing.TestUtils.createTestLlm; +import static com.google.adk.testing.TestUtils.createTextLlmResponse; import static com.google.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 java.util.stream.Collectors.joining; 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.LlmAgent; import com.google.adk.agents.NodeViewAgent; import com.google.adk.agents.RunConfig; import com.google.adk.apps.App; @@ -38,9 +43,14 @@ import com.google.adk.sessions.InMemorySessionService; import com.google.adk.sessions.Session; import com.google.adk.testing.TestBaseAgent; +import com.google.adk.testing.TestLlm; import com.google.adk.testing.TestNode; +import com.google.adk.testing.TestUtils; +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; @@ -144,6 +154,101 @@ public void runAsync_nodeRoot_persistsItsStampedEventAndState() { 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(ids(saved.immutableEvents())).containsAtLeastElementsIn(ids(events)).inOrder(); + } + + @Test + public void runAsync_agentNodesOnASubBranch_showTheLaterAgentTheEarlierAgentsReply() { + TestLlm firstModel = createTestLlm(createTextLlmResponse("from first")); + TestLlm secondModel = createTestLlm(createTextLlmResponse("from second")); + LlmAgent first = createTestAgentBuilder(firstModel).name("first").build(); + LlmAgent second = createTestAgentBuilder(secondModel).name("second").build(); + Workflow workflow = + Workflow.builder() + .name("flow") + .edges( + new Edge(Start.INSTANCE, first), + new Edge(Start.INSTANCE, TestNode.of("sibling")), + new Edge(first, second)) + .build(); + Runner runner = runner(workflow); + Session session = newSession(runner); + + runner.runAsync(USER_ID, session.id(), USER_MESSAGE).blockingSubscribe(); + + String secondRequestText = + secondModel.getLastRequest().contents().stream() + .map(TestUtils::formatContent) + .collect(joining("\n")); + assertThat(secondRequestText).contains("from first"); + } + + @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 close_workflowRoot_closesAGraphAgentThatIsAlsoASubAgentOnce() { + List closed = new ArrayList<>(); + BaseAgent helper = closeRecordingAgent("helper", closed); + BaseAgent coordinator = + new TestBaseAgent( + "coordinator", "An agent.", Flowable::empty, ImmutableList.of(helper), null, null) { + @Override + public Completable close() { + closed.add(name()); + return super.close(); + } + }; + Workflow workflow = + Workflow.builder() + .name("flow") + .edges(new Edge(Start.INSTANCE, coordinator), new Edge(coordinator, helper)) + .build(); + + runner(workflow).close().blockingAwait(); + + assertThat(closed).containsExactly("coordinator", "helper"); + } + @Test public void runAsync_failingNodeRoot_persistsItsErrorEventBeforeFailing() { IllegalStateException failure = new IllegalStateException("node failed"); @@ -535,6 +640,16 @@ private static Session reload(Runner runner, Session session) { .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", diff --git a/core/src/test/java/com/google/adk/workflow/NodeExecutionStateTest.java b/core/src/test/java/com/google/adk/workflow/NodeExecutionStateTest.java index 1463c973a..92aabdf6d 100644 --- a/core/src/test/java/com/google/adk/workflow/NodeExecutionStateTest.java +++ b/core/src/test/java/com/google/adk/workflow/NodeExecutionStateTest.java @@ -16,6 +16,7 @@ package com.google.adk.workflow; +import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertThrows; import com.google.adk.testing.TestNode; @@ -51,4 +52,17 @@ public void build_withoutARunId_throwsIllegalStateException() { assertThrows(IllegalStateException.class, builder::build); } + + @Test + public void outputFor_useAsOutputWithoutAParent_containsOnlyItsOwnPath() { + NodeExecutionState state = + NodeExecutionState.builder() + .node(TestNode.of("worker")) + .eventSink(SINK) + .runId("1") + .useAsOutput(true) + .build(); + + assertThat(state.outputFor()).containsExactly("worker@1"); + } } diff --git a/core/src/test/java/com/google/adk/workflow/NodeRunnerTest.java b/core/src/test/java/com/google/adk/workflow/NodeRunnerTest.java index 3e472dc5e..55298f358 100644 --- a/core/src/test/java/com/google/adk/workflow/NodeRunnerTest.java +++ b/core/src/test/java/com/google/adk/workflow/NodeRunnerTest.java @@ -171,6 +171,21 @@ public void run_underAParent_linksTheContextToTheParent() { 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(); @@ -764,6 +779,57 @@ public void runRoot_processingAPartialEventFails_keepsTheQueuedNodeFailureAsSupp 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 = 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..8990fa0e1 --- /dev/null +++ b/core/src/test/java/com/google/adk/workflow/SchedulerTest.java @@ -0,0 +1,676 @@ +/* + * 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 aSingleSuccessor_keepsItsPredecessorsSubBranch() { + Node a = emitter("a", "A"); + Node b = emitter("b", "B"); + + run(workflow(new Edge(START, a), new Edge(START, TestNode.of("idle")), new Edge(a, b))); + + assertThat(eventAt("wf@1/a@1").branch()).hasValue("root.a@1"); + assertThat(eventAt("wf@1/b@1").branch()).hasValue("root.a@1"); + } + + @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 disposingTheRun_cancelsTheNodesStillRunning() { + AtomicBoolean canceled = new AtomicBoolean(); + Node slow = + new TestNode( + "slow", + (ctx, input) -> + CompletableSubject.create().toFlowable().doOnCancel(() -> canceled.set(true))); + + runner(workflow(new Edge(START, slow)), "root").run(/* nodeInput= */ null).test().dispose(); + + assertThat(canceled.get()).isTrue(); + } + + @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 3e2b93e08..e4ea1318b 100644 --- a/core/src/test/java/com/google/adk/workflow/WorkflowTest.java +++ b/core/src/test/java/com/google/adk/workflow/WorkflowTest.java @@ -199,17 +199,13 @@ public void aWorkflowCanBeANodeOfAnotherWorkflow() { } @Test - public void runNode_failsAsNotSupportedYet() { - Workflow workflow = - Workflow.builder().name("flow").edges(new Edge(Start.INSTANCE, TestNode.of("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(); } }