From 5377adc7895f99247fa60f1cd9e7c9b83b6fb147 Mon Sep 17 00:00:00 2001 From: Mateusz Krawiec Date: Fri, 13 Nov 2026 10:38:36 -0800 Subject: [PATCH] fix(workflow): don't start a canceled FunctionNode's confirmation predicate PiperOrigin-RevId: 1021150226 --- .../adk/workflow/ConfirmationPredicate.java | 3 +- .../com/google/adk/workflow/FunctionNode.java | 51 ++++--- .../google/adk/workflow/FunctionNodeTest.java | 137 ++++++++++++++++-- 3 files changed, 157 insertions(+), 34 deletions(-) diff --git a/core/src/main/java/com/google/adk/workflow/ConfirmationPredicate.java b/core/src/main/java/com/google/adk/workflow/ConfirmationPredicate.java index 09218bff3..385069af4 100644 --- a/core/src/main/java/com/google/adk/workflow/ConfirmationPredicate.java +++ b/core/src/main/java/com/google/adk/workflow/ConfirmationPredicate.java @@ -23,7 +23,8 @@ /** * A predicate that decides for each input whether a {@link FunctionNode} run asks the user to * approve it, set with {@link FunctionNode.Builder#requireConfirmation(ConfirmationPredicate)}. It - * runs before the body, and an exception it throws fails the run. + * runs before the body on the body's thread. Canceling the node interrupts that thread. An + * exception it throws fails the run. */ @Experimental @FunctionalInterface diff --git a/core/src/main/java/com/google/adk/workflow/FunctionNode.java b/core/src/main/java/com/google/adk/workflow/FunctionNode.java index 6833fc453..715ff92f9 100644 --- a/core/src/main/java/com/google/adk/workflow/FunctionNode.java +++ b/core/src/main/java/com/google/adk/workflow/FunctionNode.java @@ -171,17 +171,20 @@ public Flowable runNode(Context context, @Nullable Object nodeInput) { .toFlowable(); } - /** Runs the body on the calling thread and sends its output or failure to {@code emitter}. */ + /** + * Runs the confirmation predicate and the body on the calling thread and sends the output or + * failure to {@code emitter}. Canceling the node interrupts that thread. + */ private void runBody(Context context, @Nullable Object nodeInput, MaybeEmitter emitter) { - Worker> worker = new Worker<>(() -> runFunction(context, nodeInput, emitter)); + Worker> worker = new Worker<>(() -> runFunction(context, nodeInput, emitter)); emitter.setCancellable(worker::cancel); Object output; try { - Run run = worker.run(); - if (run == null || emitter.isDisposed()) { + RunOutcome outcome = worker.run(); + if (outcome == null || emitter.isDisposed()) { return; } - output = run.replacement() != null ? run.replacement() : output(context, run); + output = outcome.replacement() != null ? outcome.replacement() : finishRun(context, outcome); } catch (Exception e) { // tryOnError, as a canceled run can no longer take a failure. emitter.tryOnError(e); @@ -198,33 +201,41 @@ private void runBody(Context context, @Nullable Object nodeInput, MaybeEmitter runFunction( + private @Nullable RunOutcome runFunction( Context context, @Nullable Object nodeInput, MaybeEmitter emitter) throws Exception { I input = inputCodec.decode(nodeInput, "input of node '" + name() + "'"); ToolConfirmation confirmation = confirmationAnswer(context); if (confirmation != null) { context.toolConfirmation(confirmation); if (!confirmation.confirmed()) { - return new Run<>(confirmation, REJECTED_OUTPUT, null); + return new RunOutcome<>(confirmation, REJECTED_OUTPUT, null); } - } else if (confirmationPredicate != null && confirmationPredicate.test(context, input)) { - return new Run<>(null, confirmationRequest(context, input), null); + } else if (confirmationPredicate != null + && !emitter.isDisposed() + && confirmationPredicate.test(context, input)) { + return new RunOutcome<>(null, confirmationRequest(context, input), null); } - return emitter.isDisposed() ? null : new Run<>(confirmation, null, body.apply(context, input)); + return emitter.isDisposed() + ? null + : new RunOutcome<>(confirmation, null, body.apply(context, input)); } - /** Returns the event or JSON value to send for a run of the body, or {@code null} for none. */ - private @Nullable Object output(Context context, Run run) { + /** + * Returns the event or JSON value to send for a run of the body, or {@code null} for none, after + * writing a delegated output in this node's output form. + */ + private @Nullable Object finishRun(Context context, RunOutcome outcome) { NodeExecutionState state = context.requireNodeState(); encodeDelegatedOutput(state); - Object output = output(run.result()); - if (run.confirmation() != null && state.delegatedOutput().isEmpty()) { - // Leaves an event for an approved run with no output, or only a dropped one, so a resume - // finds it finished. - if (output instanceof Event event && state.hasOutputDelegate()) { - return event.toBuilder().output(null).build(); + Object output = output(outcome.result()); + if (outcome.confirmation() != null && state.delegatedOutput().isEmpty()) { + // An approved run always leaves an event, so a resume finds it finished. + if (state.hasOutputDelegate()) { + return output instanceof Event event + ? event.toBuilder().output(null).build() + : finishedEvent(context); } - if (output == null || state.hasOutputDelegate()) { + if (output == null) { return finishedEvent(context); } } @@ -236,7 +247,7 @@ private void runBody(Context context, @Nullable Object nodeInput, MaybeEmitter( + private record RunOutcome( @Nullable ToolConfirmation confirmation, @Nullable Object replacement, @Nullable O result) {} /** diff --git a/core/src/test/java/com/google/adk/workflow/FunctionNodeTest.java b/core/src/test/java/com/google/adk/workflow/FunctionNodeTest.java index 71d7ce201..318da848c 100644 --- a/core/src/test/java/com/google/adk/workflow/FunctionNodeTest.java +++ b/core/src/test/java/com/google/adk/workflow/FunctionNodeTest.java @@ -122,6 +122,7 @@ public final class FunctionNodeTest { private static final String APP = "app"; private static final String USER = "user"; private static final String TRANSFER_CONFIRMATION = "wf_tool_confirmation:wf@1/transfer@1"; + private static final String PARENT_CONFIRMATION = "wf_tool_confirmation:wf@1/parent@1"; private static final String CHILD_TRANSFER_CONFIRMATION = "wf_tool_confirmation:wf@1/parent@1/transfer@1"; private static final String PAY_CONFIRMATION = "wf_tool_confirmation:wf@1/pay@1"; @@ -152,6 +153,21 @@ record Constructed(String id) { } } + /** A record whose construction blocks until interrupted, to cancel a node while it decodes. */ + record SlowToDecode(String id) { + /** Counts down when a construction starts. A test replaces it before each use. */ + static volatile CountDownLatch decoding = new CountDownLatch(1); + + SlowToDecode { + decoding.countDown(); + try { + Thread.sleep(60_000); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } + } + @JsonInclude(JsonInclude.Include.NON_NULL) record Parcel(String id, String note) {} @@ -1140,6 +1156,86 @@ public void runDisposed_sleepingBody_isInterrupted() throws InterruptedException } } + @Test + public void runDisposed_sleepingConfirmationPredicate_isInterruptedAndTheBodyDoesNotRun() + throws InterruptedException { + // Arrange: an executor scheduler, which does not interrupt on its own when disposed. + CountDownLatch started = new CountDownLatch(1); + CountDownLatch interrupted = new CountDownLatch(1); + AtomicInteger runs = new AtomicInteger(); + FunctionNode transfer = + transfer(runs) + .requireConfirmation((ctx, message) -> { + started.countDown(); + try { + Thread.sleep(60_000); + } catch (InterruptedException e) { + interrupted.countDown(); + Thread.currentThread().interrupt(); + } + return false; + }) + .build(); + ExecutorService executor = Executors.newCachedThreadPool(); + try { + Runner runner = runner(chain(transfer), Schedulers.from(executor)); + Session session = runner.sessionService().createSession(APP, USER).blockingGet(); + Disposable run = + runner.runAsync(USER, session.id(), Content.fromParts(Part.fromText("go"))).subscribe(); + assertThat(started.await(10, SECONDS)).isTrue(); + + run.dispose(); + + assertThat(interrupted.await(10, SECONDS)).isTrue(); + executor.shutdown(); + assertThat(executor.awaitTermination(10, SECONDS)).isTrue(); + assertThat(runs.get()).isEqualTo(0); + } finally { + executor.shutdownNow(); + } + } + + @Test + public void runDisposed_whileDecodingTheInput_doesNotStartTheConfirmationPredicate() + throws InterruptedException { + // Arrange: the cancel lands while the input decodes, before the predicate would start. + SlowToDecode.decoding = new CountDownLatch(1); + AtomicInteger predicateRuns = new AtomicInteger(); + AtomicInteger runs = new AtomicInteger(); + FunctionNode make = + FunctionNodes.of( + "make", Content.class, Object.class, (ctx, message) -> ImmutableMap.of("id", "a1")); + FunctionNode transfer = + FunctionNodes.builder(SlowToDecode.class, String.class) + .name("transfer") + .requireConfirmation((ctx, input) -> { + predicateRuns.incrementAndGet(); + return false; + }) + .body((ctx, input) -> { + runs.incrementAndGet(); + return "sent"; + }) + .build(); + ExecutorService executor = Executors.newCachedThreadPool(); + try { + Runner runner = runner(chain(make, transfer), Schedulers.from(executor)); + Session session = runner.sessionService().createSession(APP, USER).blockingGet(); + Disposable run = + runner.runAsync(USER, session.id(), Content.fromParts(Part.fromText("go"))).subscribe(); + assertThat(SlowToDecode.decoding.await(10, SECONDS)).isTrue(); + + run.dispose(); + + executor.shutdown(); + assertThat(executor.awaitTermination(10, SECONDS)).isTrue(); + assertThat(predicateRuns.get()).isEqualTo(0); + assertThat(runs.get()).isEqualTo(0); + } finally { + executor.shutdownNow(); + } + } + @Test public void runDisposed_bodyThatKeepsTheInterrupt_leavesTheThreadUninterruptedForTheNextTask() throws Exception { @@ -3719,6 +3815,22 @@ public void resume_approvedBodyThatWaits_rerunsOnTheAnswerWithoutAskingAgain() { assertThat(runs.get()).isEqualTo(2); } + @Test + public void resume_approvedBodyThatReturnsAnEventWithAnOutput_sendsTheOutput() { + FunctionNode transfer = + FunctionNodes.builder(Content.class, Event.class) + .name("transfer") + .requireConfirmation(true) + .body((ctx, message) -> Event.builder().output("sent").build()) + .build(); + WorkflowSession session = new WorkflowSession(chain(transfer, captureInput("reply"))); + session.advance(go()); + + ImmutableList events = session.turn(confirmation(TRANSFER_CONFIRMATION, true)); + + assertThat(replyInput(events)).isEqualTo("sent"); + } + @Test public void resume_approvedBodyThatDelegatesItsOutput_writesNoEmptyEventAndIsNotRerun() { AtomicInteger runs = new AtomicInteger(); @@ -3738,8 +3850,7 @@ public void resume_approvedBodyThatDelegatesItsOutput_writesNoEmptyEventAndIsNot .build(); WorkflowSession session = new WorkflowSession(chain(parent, new AskOnce("review", "review"))); session.advance(go()); - ImmutableList confirmed = - session.turn(confirmation("wf_tool_confirmation:wf@1/parent@1", true)); + ImmutableList confirmed = session.turn(confirmation(PARENT_CONFIRMATION, true)); session.advance(requestAnswer("review", ImmutableMap.of("ok", true))); @@ -3780,7 +3891,7 @@ public void resume_approvedBodyWhoseOnlyOutputIsDropped_isNotRerunWhenALaterNode .build(); WorkflowSession session = new WorkflowSession(chain(parent, new AskOnce("review", "review"))); session.advance(go()); - session.advance(confirmation("wf_tool_confirmation:wf@1/parent@1", true)); + session.advance(confirmation(PARENT_CONFIRMATION, true)); session.advance(requestAnswer("review", ImmutableMap.of("ok", true))); @@ -3810,8 +3921,7 @@ public void resume_approvedBodyThatRoutesAfterAnEmptyDelegation_reachesTheRouted .build()); session.advance(go()); - ImmutableList events = - session.turn(confirmation("wf_tool_confirmation:wf@1/parent@1", true)); + ImmutableList events = session.turn(confirmation(PARENT_CONFIRMATION, true)); assertThat(WorkflowResumeTesting.pathsOf(events)).contains("wf@1/reply@1"); } @@ -3823,6 +3933,7 @@ public void resume_approvedBodyThatSendsAMessageAfterAnEmptyDelegation_sendsItAn new AtomicInteger(), Event.builder() .content(Content.fromParts(Part.fromText("approved"))) + .output("dropped") .actions( EventActions.builder() .stateDelta(new HashMap<>(ImmutableMap.of("my_key", "my_value"))) @@ -3834,8 +3945,7 @@ public void resume_approvedBodyThatSendsAMessageAfterAnEmptyDelegation_sendsItAn WorkflowSession session = new WorkflowSession(chain(parent, read)); session.advance(go()); - ImmutableList events = - session.turn(confirmation("wf_tool_confirmation:wf@1/parent@1", true)); + ImmutableList events = session.turn(confirmation(PARENT_CONFIRMATION, true)); Event sent = events.stream() @@ -3854,7 +3964,7 @@ public void resume_approvedBodyThatReturnsOnlyAnOutputAfterAnEmptyDelegation_isN delegatingToASilentChild(runs, Event.builder().output("dropped").build()); WorkflowSession session = new WorkflowSession(chain(parent, new AskOnce("review", "review"))); session.advance(go()); - session.advance(confirmation("wf_tool_confirmation:wf@1/parent@1", true)); + session.advance(confirmation(PARENT_CONFIRMATION, true)); session.advance(requestAnswer("review", ImmutableMap.of("ok", true))); @@ -3868,8 +3978,7 @@ public void resume_approvedBodyThatWaitsAfterAnEmptyDelegation_asksTheUser() { WorkflowSession session = new WorkflowSession(chain(parent, captureInput("reply"))); session.advance(go()); - ImmutableList events = - session.turn(confirmation("wf_tool_confirmation:wf@1/parent@1", true)); + ImmutableList events = session.turn(confirmation(PARENT_CONFIRMATION, true)); assertThat( events.stream() @@ -4143,12 +4252,14 @@ public void toBuilder_rerunOnResume_copiesOnlyAnExplicitValue() { derived.toBuilder().requireConfirmation(false).build(); FunctionNode explicitTrueCopy = explicitTrue.toBuilder().requireConfirmation(false).build(); - FunctionNode.Builder explicitFalseCopy = - explicitFalse.toBuilder().requireConfirmation(true); + WorkflowConfigurationException explicitFalseCopyError = + assertThrows( + WorkflowConfigurationException.class, + explicitFalse.toBuilder().requireConfirmation(true)::build); assertThat(derivedCopy.rerunOnResume()).isFalse(); assertThat(explicitTrueCopy.rerunOnResume()).isTrue(); - assertThrows(WorkflowConfigurationException.class, explicitFalseCopy::build); + assertThat(explicitFalseCopyError).hasMessageThat().contains("rerunOnResume(true)"); } @Test