From 39fc4eb726d55706730605d937af29898345cbd3 Mon Sep 17 00:00:00 2001 From: Damian Momot Date: Thu, 22 Oct 2026 11:32:33 -0700 Subject: [PATCH] fix: drop a delegating node's own output, as Python ADK does PiperOrigin-RevId: 1004926780 --- .../adk/workflow/NodeExecutionState.java | 64 +++++++----- .../com/google/adk/workflow/NodeRunner.java | 42 ++++++-- .../google/adk/workflow/RunNodeConfig.java | 6 +- .../adk/workflow/DynamicDispatchTest.java | 46 +++++---- .../adk/workflow/DynamicWorkflowTest.java | 99 ++++++++++++++++--- .../adk/workflow/WorkflowResumeTest.java | 27 +++++ 6 files changed, 219 insertions(+), 65 deletions(-) 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 9c9db650a..71edc4102 100644 --- a/core/src/main/java/com/google/adk/workflow/NodeExecutionState.java +++ b/core/src/main/java/com/google/adk/workflow/NodeExecutionState.java @@ -151,26 +151,36 @@ public Optional 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(); } /** @@ -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; @@ -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); } @@ -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()); } @@ -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; } } 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 f0a1005f8..83b27351f 100644 --- a/core/src/main/java/com/google/adk/workflow/NodeRunner.java +++ b/core/src/main/java/com/google/adk/workflow/NodeRunner.java @@ -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. @@ -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 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()); } /** diff --git a/core/src/main/java/com/google/adk/workflow/RunNodeConfig.java b/core/src/main/java/com/google/adk/workflow/RunNodeConfig.java index 6287c9139..d278c6084 100644 --- a/core/src/main/java/com/google/adk/workflow/RunNodeConfig.java +++ b/core/src/main/java/com/google/adk/workflow/RunNodeConfig.java @@ -39,9 +39,9 @@ public abstract class RunNodeConfig { public abstract Optional 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(); diff --git a/core/src/test/java/com/google/adk/workflow/DynamicDispatchTest.java b/core/src/test/java/com/google/adk/workflow/DynamicDispatchTest.java index 5707955f7..4e69a477a 100644 --- a/core/src/test/java/com/google/adk/workflow/DynamicDispatchTest.java +++ b/core/src/test/java/com/google/adk/workflow/DynamicDispatchTest.java @@ -166,7 +166,7 @@ public void runChild_secondUseAsOutput_failsTheCaller() { } @Test - public void runChild_useAsOutputThenCallerProducesOutputOnItsState_failsTheCaller() { + public void runChild_useAsOutputThenCallerProducesOutputOnItsState_dropsIt() { Node dispatcher = dispatcher( "dispatcher", @@ -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 @@ -277,7 +278,7 @@ ctx, emitter("second", "SECOND"), null, USE_AS_OUTPUT) } @Test - public void runChild_delegateCanceledBeforeEmitting_keepsTheCallersOutputDelegated() { + public void runChild_delegateCanceledBeforeEmitting_keepsTheDelegationAndDropsTheFallback() { MaybeSubject cancel = MaybeSubject.create(); Node child = new TestNode( @@ -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 cancel = MaybeSubject.create(); Node child = new TestNode( @@ -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", @@ -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 diff --git a/core/src/test/java/com/google/adk/workflow/DynamicWorkflowTest.java b/core/src/test/java/com/google/adk/workflow/DynamicWorkflowTest.java index 3e9087a05..9d1337063 100644 --- a/core/src/test/java/com/google/adk/workflow/DynamicWorkflowTest.java +++ b/core/src/test/java/com/google/adk/workflow/DynamicWorkflowTest.java @@ -294,17 +294,29 @@ public void runNode_useAsOutputThroughANestedWorkflow_reachesTheCallersSuccessor } @Test - public void runNode_useAsOutputThenCallerEmitsTheChildsOutput_failsTheRun() { - Node dispatcher = returning("dispatcher", emitter("child", "CHILD_OUT"), USE_AS_OUTPUT); + public void runNode_useAsOutputThenCallerEmitsTheChildsOutput_storesOnlyTheChildsOutputEvent() { + // The caller emits the child's result, as Python ADK's delegating functions return it. + Node dispatcher = returning("dispatcher", echo("child", "echo"), USE_AS_OUTPUT); + Node after = echo("after", "after"); - RunResult result = run(workflow(new Edge(START, dispatcher))); + RunResult result = run(workflow(new Edge(START, dispatcher), new Edge(dispatcher, after))); - assertThat(result.error()).isInstanceOf(IllegalStateException.class); - assertThat(result.error()).hasMessageThat().contains("delegated its output"); + result.assertCompleted(); + ImmutableList outputEvents = result.storedOutputEvents(); + assertThat(paths(outputEvents)) + .containsExactly("wf@1/dispatcher@1/child@1", "wf@1/after@1") + .inOrder(); + assertThat(outputFor(outputEvents.get(0))) + .containsExactly("wf@1/dispatcher@1/child@1", "wf@1/dispatcher@1") + .inOrder(); + assertThat(outputEvents.get(1).output()).hasValue("after:echo:null"); + // The event that carried only the caller's own output is not sent at all. + assertThat(result.events().stream().filter(event -> pathOf(event).equals("wf@1/dispatcher@1"))) + .isEmpty(); } @Test - public void runNode_useAsOutputChildWithoutAnOutputThenCallerEmits_failsTheRun() { + public void runNode_useAsOutputChildWithoutAnOutputThenCallerEmits_dropsTheCallersOutput() { Node dispatcher = dispatcher( "dispatcher", @@ -312,18 +324,21 @@ public void runNode_useAsOutputChildWithoutAnOutputThenCallerEmits_failsTheRun() ctx.runNode(TestNode.of("child"), null, USE_AS_OUTPUT) .ignoreElement() .andThen(Flowable.just("OWN"))); + Node after = echo("after", "after"); - RunResult result = run(workflow(new Edge(START, dispatcher))); + RunResult result = run(workflow(new Edge(START, dispatcher), new Edge(dispatcher, after))); - assertThat(result.error()).isInstanceOf(IllegalStateException.class); - assertThat(result.error()).hasMessageThat().contains("delegated its output"); + result.assertCompleted(); + assertThat(result.outputs()).doesNotContain("OWN"); + assertThat(result.outputEventAt("wf@1/after@1").output()).hasValue("after:null"); } @Test - public void runNode_useAsOutputThenCallerEmitsAMessageAsOutput_failsTheRun() { + public void runNode_useAsOutputThenCallerEmitsAMessageAsOutput_storesItAsAPlainMessage() { + Content ownText = Content.fromParts(Part.fromText("own text")); Event ownMessage = newEvent() - .content(Content.fromParts(Part.fromText("own text"))) + .content(ownText) .nodeInfo(NodeInfo.builder().messageAsOutput(true).build()) .build(); Node dispatcher = @@ -333,11 +348,16 @@ public void runNode_useAsOutputThenCallerEmitsAMessageAsOutput_failsTheRun() { ctx.runNode(emitter("child", "CHILD_OUT"), null, USE_AS_OUTPUT) .ignoreElement() .andThen(Flowable.just(ownMessage))); + Node after = echo("after", "after"); - RunResult result = run(workflow(new Edge(START, dispatcher))); + RunResult result = run(workflow(new Edge(START, dispatcher), new Edge(dispatcher, after))); - assertThat(result.error()).isInstanceOf(IllegalStateException.class); - assertThat(result.error()).hasMessageThat().contains("delegated its output"); + result.assertCompleted(); + Event message = result.storedEventAt("wf@1/dispatcher@1"); + assertThat(message.content()).hasValue(ownText); + assertThat(message.nodeInfo().map(NodeInfo::messageAsOutput)).hasValue(false); + assertThat(outputFor(message)).isEmpty(); + assertThat(result.outputEventAt("wf@1/after@1").output()).hasValue("after:CHILD_OUT"); } @Test @@ -361,6 +381,57 @@ public void runNode_useAsOutputThenCallerWritesState_storesTheStateChange() { assertThat(result.eventAt("wf@1/dispatcher@1").output()).isEmpty(); } + @Test + public void runNode_useAsOutputThenCallerEmitsItsOutputWithAStateChange_storesOnlyTheChange() { + Event ownOutput = + newEvent() + .output("OWN") + .actions(EventActions.builder().stateDelta(ImmutableMap.of("written", "yes")).build()) + .build(); + Node dispatcher = + dispatcher( + "dispatcher", + (ctx, input) -> + ctx.runNode(emitter("child", "CHILD_OUT"), null, USE_AS_OUTPUT) + .ignoreElement() + .andThen(Flowable.just(ownOutput))); + Node after = echo("after", "after"); + + RunResult result = run(workflow(new Edge(START, dispatcher), new Edge(dispatcher, after))); + + result.assertCompleted(); + assertThat(result.session().state()).containsEntry("written", "yes"); + assertThat(result.storedEventAt("wf@1/dispatcher@1").output()).isEmpty(); + assertThat(result.outputEventAt("wf@1/after@1").output()).hasValue("after:CHILD_OUT"); + } + + @Test + public void runNode_useAsOutputThenCallerEmitsItsOutputWithARoute_storesOnlyTheRoute() { + Route yes = new Route.Tag("yes"); + // A resumed run reads the route back from the session. + Event ownOutput = + newEvent() + .output("OWN") + .actions(EventActions.builder().route(ImmutableList.of(yes)).build()) + .build(); + Node dispatcher = + dispatcher( + "dispatcher", + (ctx, input) -> + ctx.runNode(emitter("child", "CHILD_OUT"), null, USE_AS_OUTPUT) + .ignoreElement() + .andThen(Flowable.just(ownOutput))); + Node onYes = echo("on_yes", "yes"); + + RunResult result = run(workflow(new Edge(START, dispatcher), new Edge(dispatcher, onYes, yes))); + + result.assertCompleted(); + Event stored = result.storedEventAt("wf@1/dispatcher@1"); + assertThat(stored.actions().route()).hasValue(ImmutableList.of(yes)); + assertThat(stored.output()).isEmpty(); + assertThat(result.outputEventAt("wf@1/on_yes@1").output()).hasValue("yes:CHILD_OUT"); + } + @Test public void runNode_delegateThatFails_letsTheCallerFallBackForItsSuccessor() { Node dispatcher = diff --git a/core/src/test/java/com/google/adk/workflow/WorkflowResumeTest.java b/core/src/test/java/com/google/adk/workflow/WorkflowResumeTest.java index 16bfd6094..bb13a2ae1 100644 --- a/core/src/test/java/com/google/adk/workflow/WorkflowResumeTest.java +++ b/core/src/test/java/com/google/adk/workflow/WorkflowResumeTest.java @@ -429,6 +429,33 @@ public void resume_delegatedChildCompletedByTheAnswer_becomesTheCallersOutput() assertThat(replyInput(events)).isEqualTo(ImmutableMap.of("v", 1)); } + @Test + public void resume_callerEmittingItsReplayedDelegatesOutput_sendsThatOutputOnce() { + CountingNode child = new CountingNode("child", "child_out"); + // The caller asks after delegating, so the next turn reruns it and replays the delegate. + Node caller = + dispatcher( + "caller", + (ctx, input) -> + ctx.runNode(child, null, USE_AS_OUTPUT) + .flatMap( + result -> + ctx.runNode(longRunningAsk("ask", "ask_1"), null, runId("a")) + .map(answer -> result)) + .toFlowable()); + WorkflowSession session = + new WorkflowSession( + workflow(new Edge(START, caller), new Edge(caller, captureInput("reply")))); + session.advance(go()); + + ImmutableList events = session.turn(answer("ask_1", ImmutableMap.of("v", 1))); + + assertThat(child.runs()).isEqualTo(1); + assertThat(events.stream().filter(event -> event.errorCode().isPresent())).isEmpty(); + assertThat(outputEventAt(events, "wf@1/caller@1").output()).hasValue("child_out"); + assertThat(replyInput(events)).isEqualTo("child_out"); + } + @Test public void runNode_delegatingCallWhoseChildNoLongerRehydrates_letsTheCallerProduceItsOutput() { Node chat = new MessageOutputNode("chat", "not json", AGE_SCHEMA);