Skip to content

fix: drop a delegating node's own output, as Python ADK does #1692

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
+219 −65
Open

fix: drop a delegating node's own output, as Python ADK does #1692

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
64 changes: 42 additions & 22 deletions core/src/main/java/com/google/adk/workflow/NodeExecutionState.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -151,26 +151,36 @@ public Optional<Object> output() {

/**
* Records the node's output, which an event carries when the node finishes unless an earlier
* event already carried it. A node produces at most one output.
* event already carried it. A node produces at most one output. Once the node delegates its
* output to a {@code useAsOutput} child, an output it produces itself is dropped, as in Python
* ADK.
*
* @throws IllegalStateException if the node already has an output, or if it delegated its output
* to a {@code useAsOutput} child.
* @throws IllegalStateException if the node already produced an output of its own
*/
public void produceOutput(Object output) {
tryProduceOutput(output);
}

/**
* Records {@code output} as {@link #produceOutput} does, and returns {@code false} if it was
* dropped because a {@code useAsOutput} child holds the node's output.
*/
@CanIgnoreReturnValue
boolean tryProduceOutput(Object output) {
checkNotNull(output);
outputState.updateAndGet(
current -> {
checkState(
!current.status().isDelegated(),
"Node '%s' delegated its output to a useAsOutput child and must not produce one of"
+ " its own.",
node.name());
checkState(
current.status() == OutputStatus.NONE,
"Node '%s' produced a second output; a node produces at most one output.",
node.name());
return new OutputState(OutputStatus.PRODUCED, output);
});
OutputState updated =
outputState.updateAndGet(
current -> {
if (current.status().isDelegated()) {
return current;
}
checkState(
current.status() == OutputStatus.NONE,
"Node '%s' produced a second output; a node produces at most one output.",
node.name());
return new OutputState(OutputStatus.PRODUCED, output);
});
return !updated.status().isDelegated();
}

/**
Expand Down Expand Up @@ -288,10 +298,16 @@ void markOutputEmitted() {
.updateAndGet(
current -> {
checkState(
current.status() == OutputStatus.PRODUCED,
current.status() == OutputStatus.PRODUCED
|| current.status() == OutputStatus.DELEGATE_PRODUCED,
"Node '%s' has no unsent output to mark as emitted.",
node.name());
return new OutputState(OutputStatus.EMITTED, current.value());
// It stays delegated, so an output the node produces later is still dropped.
return new OutputState(
current.status().isDelegated()
? OutputStatus.DELEGATE_EMITTED
: OutputStatus.EMITTED,
current.value());
})
.value();
for (NodeExecutionState ancestor = outputParent;
Expand All @@ -309,7 +325,8 @@ void markOutputEmitted() {
private void recordDelegatedOutput(@Nullable Object output) {
outputState.updateAndGet(
current ->
current.status().isDelegated()
current.status() == OutputStatus.DELEGATED
|| current.status() == OutputStatus.DELEGATE_EMITTED
? new OutputState(OutputStatus.DELEGATE_EMITTED, output)
: current);
}
Expand Down Expand Up @@ -348,7 +365,8 @@ void releaseOutputDelegation() {
void adoptReplayedOutput(Object output) {
checkState(
outputState.compareAndSet(
DELEGATED_OUTPUT, new OutputState(OutputStatus.PRODUCED, checkNotNull(output))),
DELEGATED_OUTPUT,
new OutputState(OutputStatus.DELEGATE_PRODUCED, checkNotNull(output))),
"Node '%s' adopted a replayed output without a pending delegation.",
node.name());
}
Expand Down Expand Up @@ -403,11 +421,13 @@ private enum OutputStatus {
PRODUCED,
/** An event carried the output the run produced. */
EMITTED,
/** An event of a {@code useAsOutput} descendant carried the run's output. */
/** The run holds a delegate's output, such as a replayed child's, to send on its own event. */
DELEGATE_PRODUCED,
/** An event carried the run's delegated output, sent by a descendant or by the run. */
DELEGATE_EMITTED;

boolean isDelegated() {
return this == DELEGATED || this == DELEGATE_EMITTED;
return this == DELEGATED || this == DELEGATE_PRODUCED || this == DELEGATE_EMITTED;
}
}

Expand Down
42 changes: 36 additions & 6 deletions core/src/main/java/com/google/adk/workflow/NodeRunner.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -461,16 +461,24 @@ private Completable onError(Throwable error) {
}

private Completable dispatch(Event event) {
track(event);
return send(event);
boolean outputKept = track(event);
if (outputKept) {
return send(event);
}
return withoutDroppedOutput(event).map(this::send).orElseGet(Completable::complete);
}

/** Records the output, interrupts, and routing carried by {@code event} on the node state. */
private void track(Event event) {
/**
* Records the output, interrupts, and routing carried by {@code event} on the node state.
* Returns {@code false} if a {@code useAsOutput} child holds the node's output, so the event's
* own output was dropped.
*/
private boolean track(Event event) {
boolean outputKept = true;
if (event.output().isPresent()) {
nodeState.produceOutput(event.output().get());
outputKept = nodeState.tryProduceOutput(event.output().get());
} else if (isMessageAsOutput(event)) {
nodeState.produceOutput(event.content().get());
outputKept = nodeState.tryProduceOutput(event.content().get());
}
event.longRunningToolIds().ifPresent(nodeState::addInterruptIds);
// An event of a nested node was already acted on at that node's level.
Expand All @@ -485,6 +493,28 @@ private void track(Event event) {
});
event.actions().transferToAgent().ifPresent(context.eventActions()::setTransferToAgent);
}
return outputKept;
}

/**
* Returns {@code event} without the output the node produced after delegating its output, or
* empty if the event has no content, long-running tool IDs or actions to send without it. A
* message-as-output event becomes a plain message, so its text is not read back as the output.
*/
private Optional<Event> withoutDroppedOutput(Event event) {
boolean carriesMore =
event.content().isPresent()
|| event.hasPendingLongRunningToolCall()
|| !event.actions().equals(EventActions.builder().build());
if (!carriesMore) {
return Optional.empty();
}
NodeInfo plainInfo =
event
.nodeInfo()
.map(info -> info.toBuilder().messageAsOutput(false).build())
.orElse(null);
return Optional.of(event.toBuilder().output(null).nodeInfo(plainInfo).build());
}

/**
Expand Down
6 changes: 3 additions & 3 deletions core/src/main/java/com/google/adk/workflow/RunNodeConfig.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,9 +39,9 @@ public abstract class RunNodeConfig {
public abstract Optional<String> runId();

/**
* Returns whether the child's output is also the caller's output. When true, the caller must not
* produce an output of its own, even if the child produces none, unless the child fails before
* emitting its output.
* Returns whether the child's output is also the caller's output. When true, an output the caller
* produces itself is dropped, even if the child produces none. Only a child that fails before
* emitting its output gives the output back to the caller.
*/
public abstract boolean useAsOutput();

Expand Down
46 changes: 26 additions & 20 deletions core/src/test/java/com/google/adk/workflow/DynamicDispatchTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -166,7 +166,7 @@ public void runChild_secondUseAsOutput_failsTheCaller() {
}

@Test
public void runChild_useAsOutputThenCallerProducesOutputOnItsState_failsTheCaller() {
public void runChild_useAsOutputThenCallerProducesOutputOnItsState_dropsIt() {
Node dispatcher =
dispatcher(
"dispatcher",
Expand All @@ -177,10 +177,11 @@ public void runChild_useAsOutputThenCallerProducesOutputOnItsState_failsTheCalle
Completable.fromAction(() -> ctx.requireNodeState().produceOutput("OWN")))
.toFlowable());

Exception failure = runToFailure(workflow(new Edge(START, dispatcher)));
Context context = run(workflow(new Edge(START, dispatcher)));

assertThat(failure).isInstanceOf(IllegalStateException.class);
assertThat(failure).hasMessageThat().contains("delegated its output");
assertThat(context.requireNodeState().failure()).isEmpty();
assertThat(context.requireNodeState().output()).hasValue("CHILD_OUT");
assertThat(outputs()).doesNotContain("OWN");
}

@Test
Expand Down Expand Up @@ -277,7 +278,7 @@ ctx, emitter("second", "SECOND"), null, USE_AS_OUTPUT)
}

@Test
public void runChild_delegateCanceledBeforeEmitting_keepsTheCallersOutputDelegated() {
public void runChild_delegateCanceledBeforeEmitting_keepsTheDelegationAndDropsTheFallback() {
MaybeSubject<Boolean> cancel = MaybeSubject.create();
Node child =
new TestNode(
Expand All @@ -286,14 +287,16 @@ public void runChild_delegateCanceledBeforeEmitting_keepsTheCallersOutputDelegat
Completable.fromAction(() -> cancel.onSuccess(true)).andThen(Flowable.never()));
Node dispatcher = cancelingAndFallingBack("dispatcher", child, cancel);

Exception failure = runToFailure(workflow(new Edge(START, dispatcher)));
Context context = run(workflow(new Edge(START, dispatcher)));

assertThat(failure).isInstanceOf(IllegalStateException.class);
assertThat(failure).hasMessageThat().contains("delegated its output");
// Only a failed delegate releases the delegation, so the caller ends without an output.
assertThat(context.requireNodeState().failure()).isEmpty();
assertThat(context.requireNodeState().output()).isEmpty();
assertThat(outputs()).doesNotContain("FALLBACK");
}

@Test
public void runChild_delegateCanceledAfterEmittingItsOutput_failsACallerProducingItsOwn() {
public void runChild_delegateCanceledAfterEmittingItsOutput_keepsThatOutputAsTheCallers() {
MaybeSubject<Boolean> cancel = MaybeSubject.create();
Node child =
new TestNode(
Expand All @@ -304,15 +307,16 @@ public void runChild_delegateCanceledAfterEmittingItsOutput_failsACallerProducin
.concatWith(Flowable.never()));
Node dispatcher = cancelingAndFallingBack("dispatcher", child, cancel);

Exception failure = runToFailure(workflow(new Edge(START, dispatcher)));
Context context = run(workflow(new Edge(START, dispatcher)));

assertThat(failure).isInstanceOf(IllegalStateException.class);
assertThat(failure).hasMessageThat().contains("delegated its output");
assertThat(context.requireNodeState().failure()).isEmpty();
assertThat(context.requireNodeState().output()).hasValue("CHILD_OUT");
assertThat(eventAt("wf@1/dispatcher@1/child@1").output()).hasValue("CHILD_OUT");
assertThat(outputs()).doesNotContain("FALLBACK");
}

@Test
public void runChild_delegateThatFailsAfterEmittingItsOutput_failsACallerProducingItsOwn() {
public void runChild_delegateThatFailsAfterEmittingItsOutput_dropsTheCallersFallback() {
Node dispatcher =
handling(
"dispatcher",
Expand All @@ -321,22 +325,24 @@ public void runChild_delegateThatFailsAfterEmittingItsOutput_failsACallerProduci
DynamicNodeFailedException.class,
"FALLBACK");

Exception failure = runToFailure(workflow(new Edge(START, dispatcher)));
Context context = run(workflow(new Edge(START, dispatcher)));

assertThat(failure).isInstanceOf(IllegalStateException.class);
assertThat(failure).hasMessageThat().contains("delegated its output");
assertThat(context.requireNodeState().failure()).isEmpty();
assertThat(context.requireNodeState().output()).hasValue("CHILD_OUT");
assertThat(outputs()).doesNotContain("FALLBACK");
}

@Test
public void runChild_delegateWorkflowFailingAfterItsOutput_failsACallerProducingItsOwn() {
public void runChild_delegateWorkflowFailingAfterItsOutput_dropsTheCallersFallback() {
Workflow inner = workflow("inner_wf", new Edge(START, emitThenFail("term", "CHILD_OUT")));
Node dispatcher =
handling("dispatcher", inner, USE_AS_OUTPUT, DynamicNodeFailedException.class, "FALLBACK");

Exception failure = runToFailure(workflow(new Edge(START, dispatcher)));
Context context = run(workflow(new Edge(START, dispatcher)));

assertThat(failure).isInstanceOf(IllegalStateException.class);
assertThat(failure).hasMessageThat().contains("delegated its output");
assertThat(context.requireNodeState().failure()).isEmpty();
assertThat(context.requireNodeState().output()).hasValue("CHILD_OUT");
assertThat(outputs()).doesNotContain("FALLBACK");
}

@Test
Expand Down
Loading
Loading