Skip to content

Commit 145901d

Browse files
damianmomotgooglecopybara-github
authored andcommitted
refactor: add the runtime that runs child nodes from a workflow node
PiperOrigin-RevId: 1004670046
1 parent aab51ef commit 145901d

10 files changed

Lines changed: 1279 additions & 28 deletions

File tree

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
/*
2+
* Copyright 2026 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.adk.workflow;
18+
19+
import com.google.adk.annotations.Experimental;
20+
21+
/**
22+
* An exception that {@link NodeRunner#runChild} signals when the child node fails. If the caller
23+
* does not handle this exception, the caller fails with the child's error without emitting an error
24+
* event of its own.
25+
*/
26+
@Experimental
27+
final class DynamicNodeFailedException extends RuntimeException {
28+
29+
private final String errorNodePath;
30+
31+
DynamicNodeFailedException(String nodeName, Exception error, String errorNodePath) {
32+
super("Dynamic node '" + nodeName + "' failed.", error);
33+
this.errorNodePath = errorNodePath;
34+
}
35+
36+
/** Returns the exception the child failed with, which is also this exception's cause. */
37+
public Exception error() {
38+
return (Exception) getCause();
39+
}
40+
41+
/**
42+
* Returns the path of the node that raised {@link #error()}, which is either the child or one of
43+
* its descendants.
44+
*/
45+
public String errorNodePath() {
46+
return errorNodePath;
47+
}
48+
}

‎core/src/main/java/com/google/adk/workflow/NodeExecutionState.java‎

Lines changed: 141 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -25,35 +25,43 @@
2525
import com.google.common.collect.ImmutableSet;
2626
import com.google.errorprone.annotations.CanIgnoreReturnValue;
2727
import java.util.Collection;
28-
import java.util.HashMap;
29-
import java.util.LinkedHashSet;
3028
import java.util.List;
31-
import java.util.Map;
3229
import java.util.Optional;
3330
import java.util.Set;
31+
import java.util.concurrent.ConcurrentHashMap;
32+
import java.util.concurrent.ConcurrentMap;
33+
import java.util.concurrent.CopyOnWriteArraySet;
34+
import java.util.concurrent.atomic.AtomicReference;
3435
import org.jspecify.annotations.Nullable;
3536

3637
/**
3738
* The state of one workflow node run, held by the run's {@link Context}. Only the workflow runtime
38-
* creates it and updates the bookkeeping that records which results an event has carried.
39+
* creates it and updates the bookkeeping that records which results an event has carried. Its
40+
* output, interrupt IDs and child run IDs are safe to update from several threads, since children
41+
* that a node runs itself may start and finish on other threads.
3942
*/
4043
@Experimental
4144
public final class NodeExecutionState {
4245

46+
private static final OutputState NO_OUTPUT = new OutputState(OutputStatus.NONE, null);
47+
private static final OutputState DELEGATED_OUTPUT = new OutputState(OutputStatus.DELEGATED, null);
48+
4349
private final Node node;
4450
private final NodeRunner.EventSink eventSink;
4551
private final @Nullable Context parent;
52+
// The parent's state when this run's output is also the parent's.
53+
private final @Nullable NodeExecutionState outputParent;
4654
private final String runId;
4755
private final String nodePath;
4856
private final ImmutableList<String> outputFor;
49-
private final Map<String, Integer> childRunCounts = new HashMap<>();
57+
private final ConcurrentMap<String, Integer> childRunCounts = new ConcurrentHashMap<>();
58+
private final Set<String> interruptIds = new CopyOnWriteArraySet<>();
59+
private final AtomicReference<OutputState> outputState = new AtomicReference<>(NO_OUTPUT);
5060
private String eventAuthor;
51-
private final Set<String> interruptIds = new LinkedHashSet<>();
52-
private @Nullable Object output;
53-
private boolean outputEmitted;
5461
private @Nullable ImmutableList<Route> selectedRoutes;
5562
private boolean routesEmitted;
5663
private @Nullable Exception failure;
64+
private @Nullable String failureNodePath;
5765

5866
private NodeExecutionState(Builder builder) {
5967
checkState(builder.node != null, "A node state needs a node.");
@@ -64,15 +72,16 @@ private NodeExecutionState(Builder builder) {
6472
this.parent = builder.parent;
6573
this.runId = builder.runId;
6674
NodeExecutionState parentState = parent == null ? null : parent.requireNodeState();
75+
this.outputParent = builder.useAsOutput ? parentState : null;
6776
this.nodePath =
6877
NodeRunner.appendSegment(
6978
parentState == null ? null : parentState.nodePath,
7079
node.name(),
7180
runId,
7281
/* separator= */ '/');
7382
this.outputFor =
74-
builder.useAsOutput && parentState != null
75-
? ImmutableList.<String>builder().add(nodePath).addAll(parentState.outputFor).build()
83+
outputParent != null
84+
? ImmutableList.<String>builder().add(nodePath).addAll(outputParent.outputFor).build()
7685
: ImmutableList.of(nodePath);
7786
this.eventAuthor = parentState == null ? "" : parentState.eventAuthor;
7887
}
@@ -118,24 +127,36 @@ public void eventAuthor(String eventAuthor) {
118127
this.eventAuthor = checkNotNull(eventAuthor);
119128
}
120129

121-
/** Returns the node's output, or empty if it has produced none. */
130+
/**
131+
* Returns the node's output, which is the output of its {@code useAsOutput} child once that child
132+
* emits one, or empty if it has none.
133+
*/
122134
public Optional<Object> output() {
123-
return Optional.ofNullable(output);
135+
return Optional.ofNullable(outputState.get().value());
124136
}
125137

126138
/**
127139
* Records the node's output, which an event carries when the node finishes unless an earlier
128140
* event already carried it. A node produces at most one output.
129141
*
130-
* @throws IllegalStateException if the node already has an output.
142+
* @throws IllegalStateException if the node already has an output, or if it delegated its output
143+
* to a {@code useAsOutput} child.
131144
*/
132145
public void produceOutput(Object output) {
133146
checkNotNull(output);
134-
checkState(
135-
this.output == null,
136-
"Node '%s' produced a second output; a node produces at most one output.",
137-
node.name());
138-
this.output = output;
147+
outputState.updateAndGet(
148+
current -> {
149+
checkState(
150+
!current.status().isDelegated(),
151+
"Node '%s' delegated its output to a useAsOutput child and must not produce one of"
152+
+ " its own.",
153+
node.name());
154+
checkState(
155+
current.status() == OutputStatus.NONE,
156+
"Node '%s' produced a second output; a node produces at most one output.",
157+
node.name());
158+
return new OutputState(OutputStatus.PRODUCED, output);
159+
});
139160
}
140161

141162
/**
@@ -177,13 +198,73 @@ String nextChildRunId(String name) {
177198
return String.valueOf(childRunCounts.merge(name, 1, Integer::sum));
178199
}
179200

180-
/** Returns whether an event has carried the node's output. */
201+
/** Returns whether an event has carried the node's own or delegated output. */
181202
boolean outputEmitted() {
182-
return outputEmitted;
203+
OutputStatus status = outputState.get().status();
204+
return status == OutputStatus.EMITTED || status == OutputStatus.DELEGATE_EMITTED;
183205
}
184206

207+
/**
208+
* Records that an event carried the output the node produced, which also makes it the output of
209+
* each ancestor that delegated its output to this run.
210+
*
211+
* @throws IllegalStateException if the node has no unsent output.
212+
*/
185213
void markOutputEmitted() {
186-
outputEmitted = true;
214+
Object emitted =
215+
outputState
216+
.updateAndGet(
217+
current -> {
218+
checkState(
219+
current.status() == OutputStatus.PRODUCED,
220+
"Node '%s' has no unsent output to mark as emitted.",
221+
node.name());
222+
return new OutputState(OutputStatus.EMITTED, current.value());
223+
})
224+
.value();
225+
for (NodeExecutionState ancestor = outputParent;
226+
ancestor != null;
227+
ancestor = ancestor.outputParent) {
228+
ancestor.recordDelegatedOutput(emitted);
229+
}
230+
}
231+
232+
/**
233+
* Records {@code output}, emitted by a {@code useAsOutput} descendant, as the node's output if
234+
* the node delegated its output. A later output replaces an earlier one, as when a workflow
235+
* records its own output after its terminal node's.
236+
*/
237+
private void recordDelegatedOutput(@Nullable Object output) {
238+
outputState.updateAndGet(
239+
current ->
240+
current.status().isDelegated()
241+
? new OutputState(OutputStatus.DELEGATE_EMITTED, output)
242+
: current);
243+
}
244+
245+
/**
246+
* Delegates the node's output to a {@code useAsOutput} child.
247+
*
248+
* @throws IllegalStateException if the node already has an output or a delegate.
249+
*/
250+
void claimOutputDelegation() {
251+
outputState.updateAndGet(
252+
current -> {
253+
checkState(
254+
!current.status().isDelegated(),
255+
"Node '%s' already delegated its output to a useAsOutput child.",
256+
node.name());
257+
checkState(
258+
current.status() == OutputStatus.NONE,
259+
"Node '%s' already produced an output and cannot delegate it to a useAsOutput child.",
260+
node.name());
261+
return DELEGATED_OUTPUT;
262+
});
263+
}
264+
265+
/** Withdraws the delegation of the node's output, unless the delegate has emitted an output. */
266+
void releaseOutputDelegation() {
267+
outputState.compareAndSet(DELEGATED_OUTPUT, NO_OUTPUT);
187268
}
188269

189270
/** Returns whether an event has carried the node's current route selection. */
@@ -204,10 +285,49 @@ Optional<Exception> failure() {
204285
return Optional.ofNullable(failure);
205286
}
206287

288+
/** Records that the node failed with {@code failure}, which it raised itself. */
207289
void failure(Exception failure) {
290+
failure(failure, nodePath);
291+
}
292+
293+
/**
294+
* Records that the node failed with {@code failure}, which the node at {@code failureNodePath}
295+
* raised.
296+
*/
297+
void failure(Exception failure, String failureNodePath) {
208298
this.failure = checkNotNull(failure);
299+
this.failureNodePath = checkNotNull(failureNodePath);
209300
}
210301

302+
/**
303+
* Returns the path of the node that raised {@link #failure()}: this run's path, or the path of a
304+
* descendant whose failure the node carried up. Empty if the node has not failed.
305+
*/
306+
Optional<String> failureNodePath() {
307+
return Optional.ofNullable(failureNodePath);
308+
}
309+
310+
/** Status of a run's output. */
311+
private enum OutputStatus {
312+
/** The run has no output and no delegate. */
313+
NONE,
314+
/** A {@code useAsOutput} child owns the run's output and has emitted none yet. */
315+
DELEGATED,
316+
/** The run produced an output that no event has carried yet. */
317+
PRODUCED,
318+
/** An event carried the output the run produced. */
319+
EMITTED,
320+
/** An event of a {@code useAsOutput} descendant carried the run's output. */
321+
DELEGATE_EMITTED;
322+
323+
boolean isDelegated() {
324+
return this == DELEGATED || this == DELEGATE_EMITTED;
325+
}
326+
}
327+
328+
/** A run's output status and its output, held together so that each update is atomic. */
329+
private record OutputState(OutputStatus status, @Nullable Object value) {}
330+
211331
/** Builder for {@link NodeExecutionState}. */
212332
@Experimental
213333
static final class Builder {

‎core/src/main/java/com/google/adk/workflow/NodeInfo.java‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -46,9 +46,11 @@ public abstract class NodeInfo {
4646

4747
/**
4848
* Returns the node paths this event's output counts for: the emitting node's path, followed by
49-
* the paths of the ancestor workflows that use this event as their output. It is set on an event
50-
* that carries the node's output, either in its {@code output} field or as a message-as-output
51-
* event whose content is the output; an empty list means the output counts for no node.
49+
* the paths of the ancestors that use this output as their own, such as the workflow of a
50+
* terminal node or a node that ran the emitting node with {@code useAsOutput}. It is set on an
51+
* event that carries the node's output, either in its {@code output} field or as a
52+
* message-as-output event whose content is the output; an empty list means the output counts for
53+
* no node.
5254
*/
5355
@JsonProperty("outputFor")
5456
@JsonInclude(JsonInclude.Include.NON_EMPTY)
Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
/*
2+
* Copyright 2026 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.adk.workflow;
18+
19+
import com.google.adk.annotations.Experimental;
20+
21+
/**
22+
* A signal that ends a node's run when a child it ran with {@link NodeRunner#runChild} waits, so
23+
* the node waits too. Node code that catches this signal (for example, by catching every {@link
24+
* RuntimeException}) keeps running past the wait, and anything it then emits is still recorded.
25+
*/
26+
@Experimental
27+
final class NodeInterruptedException extends RuntimeException {
28+
29+
NodeInterruptedException() {
30+
super("Node interrupted.");
31+
}
32+
}

0 commit comments

Comments
 (0)