From a313ce4af4844ae36ecfbd17db3d38fc16345723 Mon Sep 17 00:00:00 2001 From: Damian Momot Date: Fri, 23 Oct 2026 06:05:08 -0700 Subject: [PATCH] feat: share one run per Context.runNode run ID PiperOrigin-RevId: 1005539735 --- .../adk/workflow/NodeExecutionState.java | 15 + .../com/google/adk/workflow/NodeRunner.java | 78 ++++- .../google/adk/workflow/RunNodeConfig.java | 7 +- .../adk/workflow/WorkflowResumeTest.java | 321 +++++++++++++++++- .../adk/workflow/WorkflowResumeTesting.java | 4 + 5 files changed, 416 insertions(+), 9 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 caf2aca6e..67a0210b5 100644 --- a/core/src/main/java/com/google/adk/workflow/NodeExecutionState.java +++ b/core/src/main/java/com/google/adk/workflow/NodeExecutionState.java @@ -26,6 +26,7 @@ import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import com.google.errorprone.annotations.CanIgnoreReturnValue; +import io.reactivex.rxjava3.subjects.SingleSubject; import java.util.Collection; import java.util.Collections; import java.util.LinkedHashMap; @@ -61,6 +62,7 @@ public final class NodeExecutionState { private final ImmutableList outputFor; private final Map resumeInputs; private final ConcurrentMap childRunCounts = new ConcurrentHashMap<>(); + private final ConcurrentMap> childRuns = new ConcurrentHashMap<>(); // The paths of the node runs started but not finished in this turn, shared under one root. private final Set unfinishedNodePaths; private final Set interruptIds = new CopyOnWriteArraySet<>(); @@ -224,6 +226,14 @@ String nextChildRunId(String name) { return String.valueOf(childRunCounts.merge(name, 1, Integer::sum)); } + /** + * Returns the runs of the children this node runs itself, keyed by {@code name@runId}, which + * calls with the same run ID share. + */ + ConcurrentMap> childRuns() { + return childRuns; + } + /** Records that this run started; it counts as canceled mid-run until it finishes. */ void markStarted() { unfinishedNodePaths.add(nodePath); @@ -255,6 +265,11 @@ synchronized ImmutableMap recoveredChildren( return recoveredChildren; } + /** Returns the parent's state if this run's output is also the parent's, or null. */ + @Nullable NodeExecutionState outputParent() { + return outputParent; + } + /** Returns whether an event has carried the node's own or delegated output. */ boolean outputEmitted() { OutputStatus status = outputState.get().status(); 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 fd9ca1e77..f0a1005f8 100644 --- a/core/src/main/java/com/google/adk/workflow/NodeRunner.java +++ b/core/src/main/java/com/google/adk/workflow/NodeRunner.java @@ -40,10 +40,12 @@ import io.reactivex.rxjava3.processors.FlowableProcessor; import io.reactivex.rxjava3.processors.UnicastProcessor; import io.reactivex.rxjava3.subjects.CompletableSubject; +import io.reactivex.rxjava3.subjects.SingleSubject; import java.util.Map; import java.util.Optional; import java.util.Queue; import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.ConcurrentMap; import java.util.function.Function; import org.jspecify.annotations.Nullable; @@ -218,8 +220,7 @@ private static Maybe startChild( parentState.claimOutputDelegation(); } String runId = config.runId().orElseGet(() -> parentState.nextChildRunId(node.name())); - // Deferred so that a throw from scan or intercept releases the delegation, as a failure does. - return Single.defer(() -> runOrReplayChild(parent, node, nodeInput, runId, config)) + return sharedChildRun(parent, node, nodeInput, runId, config) .doOnError( error -> { if (delegatesOutput) { @@ -229,6 +230,50 @@ private static Maybe startChild( .flatMapMaybe(child -> childResult(parent, child, delegatesOutput, config.raiseOnWait())); } + /** + * Returns the run of {@code node} under {@code runId} shared by the caller's calls with that run + * ID, starting it if needed. A run that fails or waits for an output it did not produce is not + * shared, so a later call runs the child again. When the call that started a run is canceled, a + * call still waiting on it runs the child itself. + */ + private static Single sharedChildRun( + Context parent, Node node, @Nullable Object nodeInput, String runId, RunNodeConfig config) { + return Single.defer( + () -> { + ConcurrentMap> childRuns = + parent.requireNodeState().childRuns(); + String childKey = segment(node.name(), runId); + SingleSubject run = SingleSubject.create(); + SingleSubject shared = childRuns.putIfAbsent(childKey, run); + if (shared != null) { + return shared.onErrorResumeNext( + error -> + error instanceof ChildRunCanceledException + ? sharedChildRun(parent, node, nodeInput, runId, config) + : Single.error(error)); + } + // Defer so a throw from scan or intercept releases the claimed run instead of leaking it. + return Single.defer(() -> runOrReplayChild(parent, node, nodeInput, runId, config)) + .doOnSuccess( + child -> { + if (!keepsItsRun(child)) { + childRuns.remove(childKey, run); + } + run.onSuccess(child); + }) + .doOnError( + error -> { + childRuns.remove(childKey, run); + run.onError(error); + }) + .doOnDispose( + () -> { + childRuns.remove(childKey, run); + run.onError(new ChildRunCanceledException()); + }); + }); + } + /** * Runs {@code node} as a child of {@code parent}, or replays the result it recorded in an earlier * turn under the same run ID. @@ -258,6 +303,22 @@ private static Single runOrReplayChild( .run(nodeInput); } + /** + * Returns whether a later call with the same run ID shares this finished run of a child: a failed + * child runs again, as does one that waits for an output and left no outcome. + */ + private static boolean keepsItsRun(Context child) { + NodeExecutionState state = child.requireNodeState(); + if (state.failure().isPresent()) { + return false; + } + return !state.interruptIds().isEmpty() + || state.output().isPresent() + || state.selectedRoutes().isPresent() + || child.eventActions().transferToAgent().isPresent() + || !state.node().waitForOutput(); + } + /** * Returns the child's output, or fails with the child's failure or with an interrupt, which the * caller takes on as its own. @@ -292,9 +353,11 @@ private static Maybe childResult( && (childNode instanceof Workflow || childNode.waitForOutput())) { return Maybe.error(new NodeInterruptedException()); } - // No event of this run carries a delegate's replayed output, so the caller emits it. + // No event of this run carries a delegate's replayed or shared output, so the caller emits it. Optional output = childState.output(); - if (delegatesOutput && output.isPresent() && !childState.outputEmitted()) { + boolean emittedForCaller = + childState.outputEmitted() && childState.outputParent() == parentState; + if (delegatesOutput && output.isPresent() && !emittedForCaller) { parentState.adoptReplayedOutput(output.get()); } return Maybe.fromOptional(output); @@ -532,6 +595,13 @@ private Event.Builder newEvent() { } } + /** Ends the waiting calls of a shared child run whose first caller was canceled. */ + private static final class ChildRunCanceledException extends RuntimeException { + ChildRunCanceledException() { + super("The call that ran this child was canceled."); + } + } + /** * An event queued for the runner, paired with the signal that completes once the caller receives * the event and, for an error event, with the failure it reports. 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 dbabfea66..6287c9139 100644 --- a/core/src/main/java/com/google/adk/workflow/RunNodeConfig.java +++ b/core/src/main/java/com/google/adk/workflow/RunNodeConfig.java @@ -31,9 +31,10 @@ public abstract class RunNodeConfig { /** * Returns the ID of the child's run, which names the run in node paths and branches. A child that - * recorded an output under the same ID in an earlier turn returns it without running again. When - * this ID is empty, the child receives the next number among the caller's runs of nodes with the - * same name, in call order, so concurrent children need explicit IDs. + * already recorded an output under the same ID, in this run of the caller or in an earlier turn, + * returns it without running again; one that finished without an output may run again. When this + * ID is empty, the child receives the next number among the caller's runs of nodes with the same + * name, in call order, so concurrent children need explicit IDs. */ public abstract Optional runId(); 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 1a32ef393..16bfd6094 100644 --- a/core/src/test/java/com/google/adk/workflow/WorkflowResumeTest.java +++ b/core/src/test/java/com/google/adk/workflow/WorkflowResumeTest.java @@ -896,6 +896,267 @@ public void resume_nodeEnteredAgainAfterTheResume_takesTheNextRunId() { assertThat(request.functionCalls().get(0).id()).hasValue("ask_2"); } + @Test + public void runNode_sameExplicitRunIdTwiceInOneRun_replaysTheFirstResult() { + CountingNode child = new CountingNode("child", "child_out"); + Node dispatcher = + dispatcher( + "dispatch", + (ctx, input) -> + ctx.runNode(child, null, runId("stable-id")) + .flatMap( + first -> + ctx.runNode(child, null, runId("stable-id")) + .map(second -> ImmutableList.of(first, second))) + .toFlowable()); + + ImmutableList events = + new WorkflowSession(workflow(new Edge(START, dispatcher))).turn(go()); + + assertThat(child.runs()).isEqualTo(1); + assertThat(outputEventAt(events, "wf@1/dispatch@1").output()) + .hasValue(ImmutableList.of("child_out", "child_out")); + } + + @Test + public void runNode_sameExplicitRunIdTwiceForAChildWithoutAnOutput_runsItOnce() { + CountingNode child = new CountingNode("child", /* output= */ null); + Node dispatcher = + dispatcher( + "dispatch", + (ctx, input) -> + ctx.runNode(child, null, runId("stable-id")) + .isEmpty() + .flatMap(unused -> ctx.runNode(child, null, runId("stable-id")).isEmpty()) + .map(secondEmpty -> "done") + .toFlowable()); + + ImmutableList events = + new WorkflowSession(workflow(new Edge(START, dispatcher))).turn(go()); + + assertThat(child.runs()).isEqualTo(1); + assertThat(outputEventAt(events, "wf@1/dispatch@1").output()).hasValue("done"); + } + + @Test + public void runNode_sameExplicitRunIdConcurrently_runsTheChildOnceForBoth() { + CompletableSubject secondJoined = CompletableSubject.create(); + CountingNode fetch = new CountingNode("fetch", "fetched", secondJoined); + Node caller = overlappingCalls("caller", fetch, runId("f"), runId("f"), secondJoined); + + ImmutableList events = new WorkflowSession(workflow(new Edge(START, caller))).turn(go()); + + assertThat(fetch.runs()).isEqualTo(1); + assertThat(outputEventAt(events, "wf@1/caller@1").output()) + .hasValue(ImmutableList.of("fetched", "fetched")); + } + + @Test + public void runNode_sameExplicitRunIdFromManyThreads_runsTheChildOnce() { + for (int i = 0; i < 20; i++) { + CountingNode fetch = new CountingNode("fetch", "fetched"); + Node caller = + dispatcher( + "caller", + (ctx, input) -> + Flowable.range(0, 8) + .flatMapMaybe( + unused -> + ctx.runNode(fetch, null, runId("f")).subscribeOn(Schedulers.io())) + .toList() + .toFlowable()); + + new WorkflowSession(workflow(new Edge(START, caller))).advance(go()); + + assertThat(fetch.runs()).isEqualTo(1); + } + } + + @Test + public void runNode_sameExplicitRunIdConcurrentlyOnAnAskingChild_asksOnce() { + AtomicInteger askRuns = new AtomicInteger(); + CompletableSubject secondJoined = CompletableSubject.create(); + Node ask = + new TestNode( + "ask", + (ctx, input) -> { + askRuns.incrementAndGet(); + return secondJoined.andThen(Flowable.just(interrupt("ask_1"))); + }); + Node caller = overlappingCalls("caller", ask, runId("f"), runId("f"), secondJoined); + + ImmutableList events = new WorkflowSession(workflow(new Edge(START, caller))).turn(go()); + + assertThat(askRuns.get()).isEqualTo(1); + assertThat(interruptEvents(events)).hasSize(1); + } + + @Test + public void runNode_sameExplicitRunIdConcurrentlyOnAFailingChild_failsBothCalls() { + AtomicInteger brokenRuns = new AtomicInteger(); + CompletableSubject secondJoined = CompletableSubject.create(); + Node broken = + new TestNode( + "broken", + (ctx, input) -> { + brokenRuns.incrementAndGet(); + return secondJoined.andThen(Flowable.error(new IllegalStateException("broken"))); + }); + Node caller = + dispatcher( + "caller", + (ctx, input) -> + Single.zip( + failedCall(ctx, broken), + failedCall(ctx, broken).doOnSubscribe(unused -> secondJoined.onComplete()), + ImmutableList::of) + .toFlowable()); + + ImmutableList events = new WorkflowSession(workflow(new Edge(START, caller))).turn(go()); + + assertThat(brokenRuns.get()).isEqualTo(1); + assertThat(outputEventAt(events, "wf@1/caller@1").output()) + .hasValue(ImmutableList.of(true, true)); + } + + @Test + public void runNode_sameExplicitRunIdAfterAFailure_runsTheChildAgain() { + AtomicInteger flakyRuns = new AtomicInteger(); + Node flaky = + new TestNode( + "flaky", + (ctx, input) -> + flakyRuns.incrementAndGet() == 1 + ? Flowable.error(new IllegalStateException("first run fails")) + : Flowable.just("recovered")); + Node caller = + dispatcher( + "caller", + (ctx, input) -> + ctx.runNode(flaky, null, runId("f")) + .onErrorResumeNext( + error -> + error instanceof DynamicNodeFailedException + ? ctx.runNode(flaky, null, runId("f")) + : Maybe.error(error)) + .toFlowable()); + + ImmutableList events = new WorkflowSession(workflow(new Edge(START, caller))).turn(go()); + + assertThat(flakyRuns.get()).isEqualTo(2); + assertThat(outputEventAt(events, "wf@1/caller@1").output()).hasValue("recovered"); + } + + @Test + public void runNode_sameExplicitRunIdAfterAWaitForOutputChildLeftNoOutput_runsTheChildAgain() { + AtomicInteger waitRuns = new AtomicInteger(); + Node wait = + new WaitingNode( + "wait", + (ctx, input) -> + waitRuns.incrementAndGet() == 1 ? Flowable.empty() : Flowable.just("ready")); + Node caller = + dispatcher( + "caller", + (ctx, input) -> + ctx.runNode(wait, null, runId("w")) + .defaultIfEmpty("none") + .flatMapMaybe( + first -> + ctx.runNode(wait, null, runId("w")) + .map(second -> ImmutableList.of(first, second))) + .toFlowable()); + + ImmutableList events = new WorkflowSession(workflow(new Edge(START, caller))).turn(go()); + + assertThat(waitRuns.get()).isEqualTo(2); + assertThat(outputEventAt(events, "wf@1/caller@1").output()) + .hasValue(ImmutableList.of("none", "ready")); + } + + @Test + public void runNode_sameExplicitRunIdAfterAnInterrupt_replaysTheWaitingChild() { + AtomicInteger askRuns = new AtomicInteger(); + Node ask = + new TestNode( + "ask", + (ctx, input) -> { + askRuns.incrementAndGet(); + return Flowable.just(interrupt("ask_1")); + }); + Node caller = + dispatcher( + "caller", + (ctx, input) -> + ctx.runNode(ask, null, runId("a")) + // Catching the wait signal keeps the caller running, as Context.runNode warns. + .onErrorResumeNext(error -> ctx.runNode(ask, null, runId("a"))) + .toFlowable()); + + ImmutableList events = new WorkflowSession(workflow(new Edge(START, caller))).turn(go()); + + assertThat(askRuns.get()).isEqualTo(1); + assertThat(interruptEvents(events)).hasSize(1); + } + + @Test + public void runNode_concurrentCallDelegatingASharedChild_emitsTheOutputForTheCaller() { + CompletableSubject secondJoined = CompletableSubject.create(); + CountingNode fetch = new CountingNode("fetch", "fetched", secondJoined); + // The first call runs fetch without delegating; the second, which delegates, shares that run. + Node caller = + dispatcher( + "caller", + (ctx, input) -> + Maybe.merge( + ctx.runNode(fetch, null, runId("f")), + ctx.runNode( + fetch, + null, + RunNodeConfig.builder().runId("f").useAsOutput(true).build()) + .doOnSubscribe(unused -> secondJoined.onComplete())) + .ignoreElements() + .toFlowable()); + + ImmutableList events = new WorkflowSession(workflow(new Edge(START, caller))).turn(go()); + + assertThat(fetch.runs()).isEqualTo(1); + assertThat(outputEventAt(events, "wf@1/caller@1").output()).hasValue("fetched"); + } + + @Test + public void runNode_firstCallOfASharedRunCanceled_waitingCallRunsTheChild() { + AtomicInteger stalledRuns = new AtomicInteger(); + CompletableSubject secondJoined = CompletableSubject.create(); + Node stalled = + new TestNode( + "stalled", + (ctx, input) -> + stalledRuns.incrementAndGet() == 1 + ? Flowable.never() + : Flowable.just("stalled done")); + Node caller = + dispatcher( + "caller", + (ctx, input) -> { + // The second call joining the run cancels the first call. + Maybe first = + ctx.runNode(stalled, null, runId("s")) + .takeUntil(secondJoined.toMaybe()) + .switchIfEmpty(Maybe.just("canceled")); + Maybe second = + ctx.runNode(stalled, null, runId("s")) + .doOnSubscribe(unused -> secondJoined.onComplete()); + return Maybe.zip(first, second, ImmutableList::of).toFlowable(); + }); + + ImmutableList events = new WorkflowSession(workflow(new Edge(START, caller))).turn(go()); + + assertThat(stalledRuns.get()).isEqualTo(2); + assertThat(outputEventAt(events, "wf@1/caller@1").output()) + .hasValue(ImmutableList.of("canceled", "stalled done")); + } + @Test public void runNode_failedChildWorkflowRunAgain_rerunsTheSiblingThatItsFailureCanceled() { AtomicInteger aRuns = new AtomicInteger(); @@ -1027,6 +1288,27 @@ public void runNode_failedCallerRunAgain_rerunsTheChildItsFailureCanceled() { assertThat(outputEventAt(events, "wf@1/caller@1").output()).hasValue("b"); } + @Test + public void runNode_childWhoseRecordedOutputNoLongerParses_failsEveryCallWithoutHanging() { + Node chat = new MessageOutputNode("chat", "not json", AGE_SCHEMA); + Node ask = longRunningAsk("ask", "ask_1"); + Node caller = + dispatcher( + "caller", + (ctx, input) -> + ctx.runNode(chat, null, runId("c")) + .onErrorResumeNext(error -> ctx.runNode(chat, null, runId("c"))) + .flatMap(unused -> ctx.runNode(ask, null, runId("a"))) + .toFlowable()); + WorkflowSession session = new WorkflowSession(workflow(new Edge(START, caller))); + session.advance(go()); + + FailedTurn turn = session.failingTurn(answer("ask_1", ImmutableMap.of("v", 1))); + + assertThat(turn.error()).isInstanceOf(IllegalArgumentException.class); + assertThat(turn.error()).hasMessageThat().contains("rehydrated output"); + } + private static ImmutableList interruptEvents(List events) { return events.stream() .filter(event -> !event.longRunningToolIds().orElse(ImmutableSet.of()).isEmpty()) @@ -1064,15 +1346,50 @@ private static Event interrupt(String interruptId) { .build(); } - /** A node that emits {@code output}, if any, and counts its runs. */ + /** + * Returns a node that runs {@code child} with {@code first} and then with {@code second}, whose + * call completes {@code secondJoined} once it has joined a run or started its own. + */ + private static Node overlappingCalls( + String name, + Node child, + RunNodeConfig first, + RunNodeConfig second, + CompletableSubject secondJoined) { + return dispatcher( + name, + (ctx, input) -> + Maybe.zip( + ctx.runNode(child, null, first), + ctx.runNode(child, null, second) + .doOnSubscribe(unused -> secondJoined.onComplete()), + ImmutableList::of) + .toFlowable()); + } + + /** Returns whether running {@code child} under run ID {@code f} fails. */ + private static Single failedCall(Context ctx, Node child) { + return ctx.runNode(child, null, runId("f")) + .map(output -> false) + .defaultIfEmpty(false) + .onErrorReturn(error -> error instanceof DynamicNodeFailedException); + } + + /** A node that emits {@code output}, if any, once {@code gate} completes, and counts its runs. */ private static final class CountingNode implements Node { private final String name; private final @Nullable Object output; + private final Completable gate; private final AtomicInteger runs = new AtomicInteger(); CountingNode(String name, @Nullable Object output) { + this(name, output, Completable.complete()); + } + + CountingNode(String name, @Nullable Object output, Completable gate) { this.name = name; this.output = output; + this.gate = gate; } int runs() { @@ -1087,7 +1404,7 @@ public String name() { @Override public Flowable runNode(Context context, @Nullable Object nodeInput) { runs.incrementAndGet(); - return output == null ? Flowable.empty() : Flowable.just(output); + return gate.andThen(output == null ? Flowable.empty() : Flowable.just(output)); } } diff --git a/core/src/test/java/com/google/adk/workflow/WorkflowResumeTesting.java b/core/src/test/java/com/google/adk/workflow/WorkflowResumeTesting.java index f6fd857e0..a19714e10 100644 --- a/core/src/test/java/com/google/adk/workflow/WorkflowResumeTesting.java +++ b/core/src/test/java/com/google/adk/workflow/WorkflowResumeTesting.java @@ -220,6 +220,10 @@ record WaitingNode( String name, BiFunction> body, boolean rerunOnResume) implements Node { + WaitingNode(String name, BiFunction> body) { + this(name, body, /* rerunOnResume= */ false); + } + @Override public boolean waitForOutput() { return true;