Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 35 additions & 0 deletions core/src/main/java/com/google/adk/runner/Runner.java
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -474,10 +477,42 @@ public PluginManager pluginManager() {
public Completable close() {
List<Completable> completables = new ArrayList<>();
completables.add(agent.close());
if (node != null) {
Set<BaseAgent> 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<BaseAgent> 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<BaseAgent> 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.
Expand Down
29 changes: 29 additions & 0 deletions core/src/main/java/com/google/adk/workflow/NodeExecutionState.java
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -43,6 +45,8 @@ public final class NodeExecutionState {
private final @Nullable Context parent;
private final String runId;
private final String nodePath;
private final ImmutableList<String> outputFor;
private final Map<String, Integer> childRunCounts = new HashMap<>();
private String eventAuthor;
private final Set<String> interruptIds = new LinkedHashSet<>();
private @Nullable Object output;
Expand All @@ -66,6 +70,10 @@ private NodeExecutionState(Builder builder) {
node.name(),
runId,
/* separator= */ '/');
this.outputFor =
builder.useAsOutput && parentState != null
? ImmutableList.<String>builder().add(nodePath).addAll(parentState.outputFor).build()
: ImmutableList.of(nodePath);
this.eventAuthor = parentState == null ? "" : parentState.eventAuthor;
}

Expand Down Expand Up @@ -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<String> 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;
Expand Down Expand Up @@ -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() {}

Expand Down Expand Up @@ -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.
*
Expand Down
18 changes: 15 additions & 3 deletions core/src/main/java/com/google/adk/workflow/NodeRunner.java
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;

Expand All @@ -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;
}
Expand Down Expand Up @@ -185,6 +186,7 @@ Single<Context> run(@Nullable Object nodeInput) {
.eventSink(eventSink)
.parent(parent)
.runId(runId)
.useAsOutput(useAsOutput)
.build());
return new Activation(context).run(nodeInput).toSingleDefault(context);
});
Expand Down Expand Up @@ -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();
}
Expand Down Expand Up @@ -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;

Expand Down Expand Up @@ -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.
Expand Down
35 changes: 35 additions & 0 deletions core/src/main/java/com/google/adk/workflow/NodeState.java
Original file line number Diff line number Diff line change
@@ -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<String> interrupts = ImmutableList.of();

NodeState(NodeStatus status) {
this.status = status;
}
}
32 changes: 32 additions & 0 deletions core/src/main/java/com/google/adk/workflow/NodeStatus.java
Original file line number Diff line number Diff line change
@@ -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
}
Loading
Loading