Skip to content

Commit 5177269

Browse files
damianmomotgooglecopybara-github
damianmomotgoogle
authored and
copybara-github
committed
fix: drop a delegating node's own output, as Python ADK does
PiperOrigin-RevId: 1007247307
1 parent eb622e6 commit 5177269

6 files changed

Lines changed: 275 additions & 63 deletions

File tree

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

Lines changed: 56 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
import java.util.concurrent.ConcurrentHashMap;
3838
import java.util.concurrent.ConcurrentMap;
3939
import java.util.concurrent.CopyOnWriteArraySet;
40+
import java.util.concurrent.atomic.AtomicBoolean;
4041
import java.util.concurrent.atomic.AtomicReference;
4142
import org.jspecify.annotations.Nullable;
4243

@@ -51,6 +52,8 @@ public final class NodeExecutionState {
5152

5253
private static final OutputState NO_OUTPUT = new OutputState(OutputStatus.NONE, null);
5354
private static final OutputState DELEGATED_OUTPUT = new OutputState(OutputStatus.DELEGATED, null);
55+
private static final String SECOND_OUTPUT_MESSAGE =
56+
"Node '%s' produced a second output; a node produces at most one output.";
5457

5558
private final Node node;
5659
private final NodeRunner.EventSink eventSink;
@@ -67,6 +70,8 @@ public final class NodeExecutionState {
6770
private final Set<String> unfinishedNodePaths;
6871
private final Set<String> interruptIds = new CopyOnWriteArraySet<>();
6972
private final AtomicReference<OutputState> outputState = new AtomicReference<>(NO_OUTPUT);
73+
// Set once the node drops its own output because a useAsOutput child holds the output.
74+
private final AtomicBoolean droppedOwnOutput = new AtomicBoolean();
7075
private String eventAuthor;
7176
private @Nullable ImmutableList<Route> selectedRoutes;
7277
private boolean routesEmitted;
@@ -151,26 +156,45 @@ public Optional<Object> output() {
151156

152157
/**
153158
* Records the node's output, which an event carries when the node finishes unless an earlier
154-
* event already carried it. A node produces at most one output.
159+
* event already carried it. A node produces at most one output. Once the node delegates its
160+
* output to a {@code useAsOutput} child, an output it produces itself is dropped, as in Python
161+
* ADK.
155162
*
156-
* @throws IllegalStateException if the node already has an output, or if it delegated its output
157-
* to a {@code useAsOutput} child.
163+
* @throws IllegalStateException if the node already produced an output of its own, even one that
164+
* was dropped
158165
*/
159166
public void produceOutput(Object output) {
167+
tryProduceOutput(output);
168+
}
169+
170+
/**
171+
* Records {@code output} as {@link #produceOutput} does, and returns {@code false} if it was
172+
* dropped because a {@code useAsOutput} child holds the node's output.
173+
*
174+
* @throws IllegalStateException if the node already produced an output of its own, even one that
175+
* was dropped
176+
*/
177+
@CanIgnoreReturnValue
178+
boolean tryProduceOutput(Object output) {
160179
checkNotNull(output);
161-
outputState.updateAndGet(
162-
current -> {
163-
checkState(
164-
!current.status().isDelegated(),
165-
"Node '%s' delegated its output to a useAsOutput child and must not produce one of"
166-
+ " its own.",
167-
node.name());
168-
checkState(
169-
current.status() == OutputStatus.NONE,
170-
"Node '%s' produced a second output; a node produces at most one output.",
171-
node.name());
172-
return new OutputState(OutputStatus.PRODUCED, output);
173-
});
180+
OutputState updated =
181+
outputState.updateAndGet(
182+
current -> {
183+
if (current.status().isDelegated()) {
184+
return current;
185+
}
186+
checkState(
187+
current.status() == OutputStatus.NONE && !droppedOwnOutput.get(),
188+
SECOND_OUTPUT_MESSAGE,
189+
node.name());
190+
return new OutputState(OutputStatus.PRODUCED, output);
191+
});
192+
if (!updated.status().isDelegated()) {
193+
return true;
194+
}
195+
// As in Python ADK, a dropped output still counts, so a later one fails as a second output.
196+
checkState(droppedOwnOutput.compareAndSet(false, true), SECOND_OUTPUT_MESSAGE, node.name());
197+
return false;
174198
}
175199

176200
/**
@@ -288,10 +312,16 @@ void markOutputEmitted() {
288312
.updateAndGet(
289313
current -> {
290314
checkState(
291-
current.status() == OutputStatus.PRODUCED,
315+
current.status() == OutputStatus.PRODUCED
316+
|| current.status() == OutputStatus.DELEGATE_PRODUCED,
292317
"Node '%s' has no unsent output to mark as emitted.",
293318
node.name());
294-
return new OutputState(OutputStatus.EMITTED, current.value());
319+
// It stays delegated, so an output the node produces later is still dropped.
320+
return new OutputState(
321+
current.status().isDelegated()
322+
? OutputStatus.DELEGATE_EMITTED
323+
: OutputStatus.EMITTED,
324+
current.value());
295325
})
296326
.value();
297327
for (NodeExecutionState ancestor = outputParent;
@@ -309,7 +339,8 @@ void markOutputEmitted() {
309339
private void recordDelegatedOutput(@Nullable Object output) {
310340
outputState.updateAndGet(
311341
current ->
312-
current.status().isDelegated()
342+
current.status() == OutputStatus.DELEGATED
343+
|| current.status() == OutputStatus.DELEGATE_EMITTED
313344
? new OutputState(OutputStatus.DELEGATE_EMITTED, output)
314345
: current);
315346
}
@@ -348,7 +379,8 @@ void releaseOutputDelegation() {
348379
void adoptReplayedOutput(Object output) {
349380
checkState(
350381
outputState.compareAndSet(
351-
DELEGATED_OUTPUT, new OutputState(OutputStatus.PRODUCED, checkNotNull(output))),
382+
DELEGATED_OUTPUT,
383+
new OutputState(OutputStatus.DELEGATE_PRODUCED, checkNotNull(output))),
352384
"Node '%s' adopted a replayed output without a pending delegation.",
353385
node.name());
354386
}
@@ -403,11 +435,13 @@ private enum OutputStatus {
403435
PRODUCED,
404436
/** An event carried the output the run produced. */
405437
EMITTED,
406-
/** An event of a {@code useAsOutput} descendant carried the run's output. */
438+
/** The run holds a delegate's output, such as a replayed child's, to send on its own event. */
439+
DELEGATE_PRODUCED,
440+
/** An event carried the run's delegated output, sent by a descendant or by the run. */
407441
DELEGATE_EMITTED;
408442

409443
boolean isDelegated() {
410-
return this == DELEGATED || this == DELEGATE_EMITTED;
444+
return this == DELEGATED || this == DELEGATE_PRODUCED || this == DELEGATE_EMITTED;
411445
}
412446
}
413447

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

Lines changed: 36 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -461,16 +461,24 @@ private Completable onError(Throwable error) {
461461
}
462462

463463
private Completable dispatch(Event event) {
464-
track(event);
465-
return send(event);
464+
boolean outputKept = track(event);
465+
if (outputKept) {
466+
return send(event);
467+
}
468+
return withoutDroppedOutput(event).map(this::send).orElseGet(Completable::complete);
466469
}
467470

468-
/** Records the output, interrupts, and routing carried by {@code event} on the node state. */
469-
private void track(Event event) {
471+
/**
472+
* Records the output, interrupts, and routing carried by {@code event} on the node state.
473+
* Returns {@code false} if a {@code useAsOutput} child holds the node's output, so the event's
474+
* own output was dropped.
475+
*/
476+
private boolean track(Event event) {
477+
boolean outputKept = true;
470478
if (event.output().isPresent()) {
471-
nodeState.produceOutput(event.output().get());
479+
outputKept = nodeState.tryProduceOutput(event.output().get());
472480
} else if (isMessageAsOutput(event)) {
473-
nodeState.produceOutput(event.content().get());
481+
outputKept = nodeState.tryProduceOutput(event.content().get());
474482
}
475483
event.longRunningToolIds().ifPresent(nodeState::addInterruptIds);
476484
// An event of a nested node was already acted on at that node's level.
@@ -485,6 +493,28 @@ private void track(Event event) {
485493
});
486494
event.actions().transferToAgent().ifPresent(context.eventActions()::setTransferToAgent);
487495
}
496+
return outputKept;
497+
}
498+
499+
/**
500+
* Returns {@code event} without the output the node produced after delegating its output, or
501+
* empty if the event has no content, long-running tool IDs or actions to send without it. A
502+
* message-as-output event becomes a plain message, so its text is not read back as the output.
503+
*/
504+
private Optional<Event> withoutDroppedOutput(Event event) {
505+
boolean carriesMore =
506+
event.content().isPresent()
507+
|| event.hasPendingLongRunningToolCall()
508+
|| !event.actions().equals(EventActions.builder().build());
509+
if (!carriesMore) {
510+
return Optional.empty();
511+
}
512+
NodeInfo plainInfo =
513+
event
514+
.nodeInfo()
515+
.map(info -> info.toBuilder().messageAsOutput(false).build())
516+
.orElse(null);
517+
return Optional.of(event.toBuilder().output(null).nodeInfo(plainInfo).build());
488518
}
489519

490520
/**

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -39,9 +39,9 @@ public abstract class RunNodeConfig {
3939
public abstract Optional<String> runId();
4040

4141
/**
42-
* Returns whether the child's output is also the caller's output. When true, the caller must not
43-
* produce an output of its own, even if the child produces none, unless the child fails before
44-
* emitting its output.
42+
* Returns whether the child's output is also the caller's output. When true, an output the caller
43+
* produces itself is dropped even if the child produces none, and a second one still fails the
44+
* caller. Only a child that fails before emitting its output gives the output back to the caller.
4545
*/
4646
public abstract boolean useAsOutput();
4747

‎core/src/test/java/com/google/adk/workflow/DynamicDispatchTest.java‎

Lines changed: 68 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -166,7 +166,7 @@ public void runChild_secondUseAsOutput_failsTheCaller() {
166166
}
167167

168168
@Test
169-
public void runChild_useAsOutputThenCallerProducesOutputOnItsState_failsTheCaller() {
169+
public void runChild_useAsOutputThenCallerProducesOutputOnItsState_dropsIt() {
170170
Node dispatcher =
171171
dispatcher(
172172
"dispatcher",
@@ -177,10 +177,32 @@ public void runChild_useAsOutputThenCallerProducesOutputOnItsState_failsTheCalle
177177
Completable.fromAction(() -> ctx.requireNodeState().produceOutput("OWN")))
178178
.toFlowable());
179179

180+
Context context = run(workflow(new Edge(START, dispatcher)));
181+
182+
assertThat(context.requireNodeState().failure()).isEmpty();
183+
assertThat(context.requireNodeState().output()).hasValue("CHILD_OUT");
184+
assertThat(outputs()).doesNotContain("OWN");
185+
}
186+
187+
@Test
188+
public void runChild_useAsOutputThenCallerProducesTwoOutputs_failsTheCaller() {
189+
Node dispatcher =
190+
dispatcher(
191+
"dispatcher",
192+
(ctx, input) ->
193+
NodeRunner.runChild(ctx, emitter("child", "CHILD_OUT"), null, USE_AS_OUTPUT)
194+
.ignoreElement()
195+
.doOnComplete(
196+
() -> {
197+
ctx.requireNodeState().produceOutput("OWN");
198+
ctx.requireNodeState().produceOutput("SECOND");
199+
})
200+
.toFlowable());
201+
180202
Exception failure = runToFailure(workflow(new Edge(START, dispatcher)));
181203

182204
assertThat(failure).isInstanceOf(IllegalStateException.class);
183-
assertThat(failure).hasMessageThat().contains("delegated its output");
205+
assertThat(failure).hasMessageThat().contains("produced a second output");
184206
}
185207

186208
@Test
@@ -249,6 +271,29 @@ public void runChild_delegateThatFailsBeforeEmittingItsOutput_letsTheCallerProdu
249271
assertThat(outputs()).doesNotContain("CHILD_OUT");
250272
}
251273

274+
@Test
275+
public void runChild_callerOutputDroppedThenDelegateFailsBeforeEmitting_failsTheFallback() {
276+
Node child =
277+
new TestNode(
278+
"child",
279+
(ctx, input) -> {
280+
// Simulates the caller producing an output while its delegate still runs.
281+
ctx.requireNodeState()
282+
.parent()
283+
.orElseThrow()
284+
.requireNodeState()
285+
.produceOutput("EARLY");
286+
return Flowable.error(new IllegalStateException("failed before emitting"));
287+
});
288+
Node dispatcher =
289+
handling("dispatcher", child, USE_AS_OUTPUT, DynamicNodeFailedException.class, "FALLBACK");
290+
291+
Exception failure = runToFailure(workflow(new Edge(START, dispatcher)));
292+
293+
assertThat(failure).isInstanceOf(IllegalStateException.class);
294+
assertThat(failure).hasMessageThat().contains("produced a second output");
295+
}
296+
252297
@Test
253298
public void runChild_delegateThatFailsBeforeEmitting_letsTheCallerDelegateAgain() {
254299
Node dispatcher =
@@ -277,7 +322,7 @@ ctx, emitter("second", "SECOND"), null, USE_AS_OUTPUT)
277322
}
278323

279324
@Test
280-
public void runChild_delegateCanceledBeforeEmitting_keepsTheCallersOutputDelegated() {
325+
public void runChild_delegateCanceledBeforeEmitting_keepsTheDelegationAndDropsTheFallback() {
281326
MaybeSubject<Boolean> cancel = MaybeSubject.create();
282327
Node child =
283328
new TestNode(
@@ -286,14 +331,16 @@ public void runChild_delegateCanceledBeforeEmitting_keepsTheCallersOutputDelegat
286331
Completable.fromAction(() -> cancel.onSuccess(true)).andThen(Flowable.never()));
287332
Node dispatcher = cancelingAndFallingBack("dispatcher", child, cancel);
288333

289-
Exception failure = runToFailure(workflow(new Edge(START, dispatcher)));
334+
Context context = run(workflow(new Edge(START, dispatcher)));
290335

291-
assertThat(failure).isInstanceOf(IllegalStateException.class);
292-
assertThat(failure).hasMessageThat().contains("delegated its output");
336+
// Only a failed delegate releases the delegation, so the caller ends without an output.
337+
assertThat(context.requireNodeState().failure()).isEmpty();
338+
assertThat(context.requireNodeState().output()).isEmpty();
339+
assertThat(outputs()).doesNotContain("FALLBACK");
293340
}
294341

295342
@Test
296-
public void runChild_delegateCanceledAfterEmittingItsOutput_failsACallerProducingItsOwn() {
343+
public void runChild_delegateCanceledAfterEmittingItsOutput_keepsThatOutputAsTheCallers() {
297344
MaybeSubject<Boolean> cancel = MaybeSubject.create();
298345
Node child =
299346
new TestNode(
@@ -304,15 +351,16 @@ public void runChild_delegateCanceledAfterEmittingItsOutput_failsACallerProducin
304351
.concatWith(Flowable.never()));
305352
Node dispatcher = cancelingAndFallingBack("dispatcher", child, cancel);
306353

307-
Exception failure = runToFailure(workflow(new Edge(START, dispatcher)));
354+
Context context = run(workflow(new Edge(START, dispatcher)));
308355

309-
assertThat(failure).isInstanceOf(IllegalStateException.class);
310-
assertThat(failure).hasMessageThat().contains("delegated its output");
356+
assertThat(context.requireNodeState().failure()).isEmpty();
357+
assertThat(context.requireNodeState().output()).hasValue("CHILD_OUT");
311358
assertThat(eventAt("wf@1/dispatcher@1/child@1").output()).hasValue("CHILD_OUT");
359+
assertThat(outputs()).doesNotContain("FALLBACK");
312360
}
313361

314362
@Test
315-
public void runChild_delegateThatFailsAfterEmittingItsOutput_failsACallerProducingItsOwn() {
363+
public void runChild_delegateThatFailsAfterEmittingItsOutput_dropsTheCallersFallback() {
316364
Node dispatcher =
317365
handling(
318366
"dispatcher",
@@ -321,22 +369,24 @@ public void runChild_delegateThatFailsAfterEmittingItsOutput_failsACallerProduci
321369
DynamicNodeFailedException.class,
322370
"FALLBACK");
323371

324-
Exception failure = runToFailure(workflow(new Edge(START, dispatcher)));
372+
Context context = run(workflow(new Edge(START, dispatcher)));
325373

326-
assertThat(failure).isInstanceOf(IllegalStateException.class);
327-
assertThat(failure).hasMessageThat().contains("delegated its output");
374+
assertThat(context.requireNodeState().failure()).isEmpty();
375+
assertThat(context.requireNodeState().output()).hasValue("CHILD_OUT");
376+
assertThat(outputs()).doesNotContain("FALLBACK");
328377
}
329378

330379
@Test
331-
public void runChild_delegateWorkflowFailingAfterItsOutput_failsACallerProducingItsOwn() {
380+
public void runChild_delegateWorkflowFailingAfterItsOutput_dropsTheCallersFallback() {
332381
Workflow inner = workflow("inner_wf", new Edge(START, emitThenFail("term", "CHILD_OUT")));
333382
Node dispatcher =
334383
handling("dispatcher", inner, USE_AS_OUTPUT, DynamicNodeFailedException.class, "FALLBACK");
335384

336-
Exception failure = runToFailure(workflow(new Edge(START, dispatcher)));
385+
Context context = run(workflow(new Edge(START, dispatcher)));
337386

338-
assertThat(failure).isInstanceOf(IllegalStateException.class);
339-
assertThat(failure).hasMessageThat().contains("delegated its output");
387+
assertThat(context.requireNodeState().failure()).isEmpty();
388+
assertThat(context.requireNodeState().output()).hasValue("CHILD_OUT");
389+
assertThat(outputs()).doesNotContain("FALLBACK");
340390
}
341391

342392
@Test

0 commit comments

Comments
 (0)