Skip to content

fix(workflow): don't start a canceled FunctionNode's confirmation predicate #1806

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
+157 −34
Open

fix(workflow): don't start a canceled FunctionNode's confirmation predicate #1806

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
3 changes: 2 additions & 1 deletion core/src/main/java/com/google/adk/workflow/ConfirmationPredicate.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
51 changes: 31 additions & 20 deletions core/src/main/java/com/google/adk/workflow/FunctionNode.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<Object> emitter) {
Worker<Run<O>> worker = new Worker<>(() -> runFunction(context, nodeInput, emitter));
Worker<RunOutcome<O>> worker = new Worker<>(() -> runFunction(context, nodeInput, emitter));
emitter.setCancellable(worker::cancel);
Object output;
try {
Run<O> run = worker.run();
if (run == null || emitter.isDisposed()) {
RunOutcome<O> 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);
Expand All @@ -198,33 +201,41 @@ private void runBody(Context context, @Nullable Object nodeInput, MaybeEmitter<O
* Asks for or reads the user's confirmation, then runs the body unless the rejection or the
* request replaces it. Runs in the worker, as the predicate is user code like the body.
*/
private @Nullable Run<O> runFunction(
private @Nullable RunOutcome<O> runFunction(
Context context, @Nullable Object nodeInput, MaybeEmitter<Object> 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<O> 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<O> 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);
}
}
Expand All @@ -236,7 +247,7 @@ private void runBody(Context context, @Nullable Object nodeInput, MaybeEmitter<O
* One run: the user's answer to its confirmation request, if any, and either what replaces the
* body (a rejection or a confirmation request) or the body's result.
*/
private record Run<O>(
private record RunOutcome<O>(
@Nullable ToolConfirmation confirmation, @Nullable Object replacement, @Nullable O result) {}

/**
Expand Down
137 changes: 124 additions & 13 deletions core/src/test/java/com/google/adk/workflow/FunctionNodeTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -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) {}

Expand Down Expand Up @@ -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<Content, String> 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<Content, Object> make =
FunctionNodes.of(
"make", Content.class, Object.class, (ctx, message) -> ImmutableMap.of("id", "a1"));
FunctionNode<SlowToDecode, String> 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 {
Expand Down Expand Up @@ -3719,6 +3815,22 @@ public void resume_approvedBodyThatWaits_rerunsOnTheAnswerWithoutAskingAgain() {
assertThat(runs.get()).isEqualTo(2);
}

@Test
public void resume_approvedBodyThatReturnsAnEventWithAnOutput_sendsTheOutput() {
FunctionNode<Content, Event> 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<Event> events = session.turn(confirmation(TRANSFER_CONFIRMATION, true));

assertThat(replyInput(events)).isEqualTo("sent");
}

@Test
public void resume_approvedBodyThatDelegatesItsOutput_writesNoEmptyEventAndIsNotRerun() {
AtomicInteger runs = new AtomicInteger();
Expand All @@ -3738,8 +3850,7 @@ public void resume_approvedBodyThatDelegatesItsOutput_writesNoEmptyEventAndIsNot
.build();
WorkflowSession session = new WorkflowSession(chain(parent, new AskOnce("review", "review")));
session.advance(go());
ImmutableList<Event> confirmed =
session.turn(confirmation("wf_tool_confirmation:wf@1/parent@1", true));
ImmutableList<Event> confirmed = session.turn(confirmation(PARENT_CONFIRMATION, true));

session.advance(requestAnswer("review", ImmutableMap.of("ok", true)));

Expand Down Expand Up @@ -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)));

Expand Down Expand Up @@ -3810,8 +3921,7 @@ public void resume_approvedBodyThatRoutesAfterAnEmptyDelegation_reachesTheRouted
.build());
session.advance(go());

ImmutableList<Event> events =
session.turn(confirmation("wf_tool_confirmation:wf@1/parent@1", true));
ImmutableList<Event> events = session.turn(confirmation(PARENT_CONFIRMATION, true));

assertThat(WorkflowResumeTesting.pathsOf(events)).contains("wf@1/reply@1");
}
Expand All @@ -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")))
Expand All @@ -3834,8 +3945,7 @@ public void resume_approvedBodyThatSendsAMessageAfterAnEmptyDelegation_sendsItAn
WorkflowSession session = new WorkflowSession(chain(parent, read));
session.advance(go());

ImmutableList<Event> events =
session.turn(confirmation("wf_tool_confirmation:wf@1/parent@1", true));
ImmutableList<Event> events = session.turn(confirmation(PARENT_CONFIRMATION, true));

Event sent =
events.stream()
Expand All @@ -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)));

Expand All @@ -3868,8 +3978,7 @@ public void resume_approvedBodyThatWaitsAfterAnEmptyDelegation_asksTheUser() {
WorkflowSession session = new WorkflowSession(chain(parent, captureInput("reply")));
session.advance(go());

ImmutableList<Event> events =
session.turn(confirmation("wf_tool_confirmation:wf@1/parent@1", true));
ImmutableList<Event> events = session.turn(confirmation(PARENT_CONFIRMATION, true));

assertThat(
events.stream()
Expand Down Expand Up @@ -4143,12 +4252,14 @@ public void toBuilder_rerunOnResume_copiesOnlyAnExplicitValue() {
derived.toBuilder().requireConfirmation(false).build();
FunctionNode<Content, String> explicitTrueCopy =
explicitTrue.toBuilder().requireConfirmation(false).build();
FunctionNode.Builder<Content, String> 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
Expand Down
Loading