From 1b112908408a0a1219a40cbb86264db58478e835 Mon Sep 17 00:00:00 2001 From: Mateusz Krawiec Date: Fri, 13 Nov 2026 08:22:55 -0800 Subject: [PATCH] feat(workflow): ask the user to approve a FunctionNode run PiperOrigin-RevId: 1021052780 --- .../adk/workflow/ConfirmationPredicate.java | 37 + .../com/google/adk/workflow/FunctionNode.java | 206 ++++- .../adk/workflow/NodeExecutionState.java | 5 + .../google/adk/workflow/FunctionNodeTest.java | 823 ++++++++++++++++++ 4 files changed, 1054 insertions(+), 17 deletions(-) create mode 100644 core/src/main/java/com/google/adk/workflow/ConfirmationPredicate.java diff --git a/core/src/main/java/com/google/adk/workflow/ConfirmationPredicate.java b/core/src/main/java/com/google/adk/workflow/ConfirmationPredicate.java new file mode 100644 index 000000000..09218bff3 --- /dev/null +++ b/core/src/main/java/com/google/adk/workflow/ConfirmationPredicate.java @@ -0,0 +1,37 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.adk.workflow; + +import com.google.adk.agents.Context; +import com.google.adk.annotations.Experimental; +import org.jspecify.annotations.Nullable; + +/** + * 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. + */ +@Experimental +@FunctionalInterface +public interface ConfirmationPredicate { + + /** + * Returns whether this run asks the user to approve it before the body runs. {@code input} is + * null when the predecessor sent no output or the port is {@code Void}. + */ + boolean test(Context context, @Nullable I input) throws Exception; +} 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 70e79b9d9..6833fc453 100644 --- a/core/src/main/java/com/google/adk/workflow/FunctionNode.java +++ b/core/src/main/java/com/google/adk/workflow/FunctionNode.java @@ -23,12 +23,17 @@ import com.google.adk.agents.Context; import com.google.adk.annotations.Experimental; import com.google.adk.events.Event; +import com.google.adk.events.ToolConfirmation; import com.google.common.base.Throwables; +import com.google.common.collect.ImmutableMap; import com.google.errorprone.annotations.CanIgnoreReturnValue; import com.google.genai.types.Schema; +import com.google.genai.types.Type; import io.reactivex.rxjava3.core.Flowable; import io.reactivex.rxjava3.core.Maybe; import io.reactivex.rxjava3.core.MaybeEmitter; +import java.util.LinkedHashMap; +import java.util.Map; import java.util.Optional; import java.util.concurrent.Callable; import java.util.concurrent.ExecutionException; @@ -51,15 +56,38 @@ public final class FunctionNode extends BaseNode { private static final PortCodec ANY = PortCodec.output(PortCodec.javaType(Object.class), "an event output"); + private static final String CONFIRMATION_INTERRUPT_PREFIX = "wf_tool_confirmation:"; + + private static final ConfirmationPredicate ALWAYS_CONFIRM = (context, input) -> true; + + /** The output of a run the user rejected. */ + private static final ImmutableMap REJECTED_OUTPUT = + ImmutableMap.of("error", "This tool call is rejected."); + + /** The schema of an answer to a confirmation request. */ + private static final Schema TOOL_CONFIRMATION_SCHEMA = + Schema.builder() + .type(Type.Known.OBJECT) + .properties( + ImmutableMap.of( + "hint", Schema.builder().type(Type.Known.STRING).nullable(true).build(), + "confirmed", Schema.builder().type(Type.Known.BOOLEAN).nullable(true).build(), + "payload", Schema.builder().build())) + .build(); + private final PortCodec inputCodec; private final PortCodec outputCodec; private final NodeBody body; + private final @Nullable ConfirmationPredicate confirmationPredicate; + private final @Nullable Boolean explicitRerunOnResume; private FunctionNode(Builder builder) { super( checkNotNull(builder.name), builder.description, - builder.rerunOnResume, + builder.rerunOnResume != null + ? builder.rerunOnResume + : builder.confirmationPredicate != null, builder.waitForOutput, builder.config, builder.inputSchema, @@ -68,14 +96,21 @@ private FunctionNode(Builder builder) { this.inputCodec = builder.inputCodec; this.outputCodec = builder.outputCodec; this.body = checkNotNull(builder.body); + this.confirmationPredicate = builder.confirmationPredicate; + this.explicitRerunOnResume = builder.rerunOnResume; } - /** Returns a builder initialized with this node's properties. */ + /** + * Returns a builder initialized with this node's properties. A {@code rerunOnResume} the node + * derived from its confirmation is left unset, so the copy derives it from its own. + */ public Builder toBuilder() { - return new Builder<>(inputCodec, outputCodec) + Builder builder = new Builder<>(inputCodec, outputCodec); + builder.rerunOnResume = explicitRerunOnResume; + builder.confirmationPredicate = confirmationPredicate; + return builder .name(name()) .description(description()) - .rerunOnResume(rerunOnResume()) .waitForOutput(waitForOutput()) .config(config()) // Only explicit schemas carry over; the copy infers the others from its ports. @@ -138,21 +173,15 @@ public Flowable runNode(Context context, @Nullable Object nodeInput) { /** Runs the body on the calling thread and sends its output or failure to {@code emitter}. */ private void runBody(Context context, @Nullable Object nodeInput, MaybeEmitter emitter) { - Worker worker = - new Worker<>(() -> { - I input = inputCodec.decode(nodeInput, "input of node '" + name() + "'"); - return emitter.isDisposed() ? null : body.apply(context, input); - }); + Worker> worker = new Worker<>(() -> runFunction(context, nodeInput, emitter)); emitter.setCancellable(worker::cancel); Object output; try { - O result = worker.run(); - if (emitter.isDisposed()) { + Run run = worker.run(); + if (run == null || emitter.isDisposed()) { return; } - encodeDelegatedOutput(context.requireNodeState()); - // After a delegation, NodeRunner drops a returned output. - output = output(result); + output = run.replacement() != null ? run.replacement() : output(context, run); } catch (Exception e) { // tryOnError, as a canceled run can no longer take a failure. emitter.tryOnError(e); @@ -165,6 +194,51 @@ private void runBody(Context context, @Nullable Object nodeInput, MaybeEmitter 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); + } + } else if (confirmationPredicate != null && confirmationPredicate.test(context, input)) { + return new Run<>(null, confirmationRequest(context, input), null); + } + return emitter.isDisposed() ? null : new Run<>(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) { + 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(); + } + if (output == null || state.hasOutputDelegate()) { + return finishedEvent(context); + } + } + // After a delegation, NodeRunner drops a returned output. + return output; + } + + /** + * 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( + @Nullable ToolConfirmation confirmation, @Nullable Object replacement, @Nullable O result) {} + /** * Writes the output of a child the body gave its output to with {@code useAsOutput} as if the * body had returned it, and has the node send it as its own output if that changes it. @@ -199,6 +273,65 @@ private void encodeDelegatedOutput(NodeExecutionState state) { } } + private static Event finishedEvent(Context context) { + return Event.builder() + .id(context.invocationContext().newUuid()) + .timestamp(context.invocationContext().now().toEpochMilli()) + .author("") + .build(); + } + + private static String confirmationInterruptId(Context context) { + return CONFIRMATION_INTERRUPT_PREFIX + context.requireNodeState().nodePath(); + } + + /** Returns the request that asks the user to approve this run with {@code input}. */ + private RequestInput confirmationRequest(Context context, @Nullable I input) { + String interruptId = confirmationInterruptId(context); + Map payload = new LinkedHashMap<>(); + payload.put("tool_name", name()); + payload.put("args", inputCodec.encode(input, "input of node '" + name() + "'")); + return RequestInput.builder() + .interruptId(interruptId) + .message("Please approve or reject the tool call " + name() + "().") + .payload(payload) + .responseSchema(TOOL_CONFIRMATION_SCHEMA) + .build(); + } + + /** + * Returns the user's answer to this run's confirmation request, or null if the node requires no + * confirmation or has no answer yet. A field left out or null takes its default. + * + * @throws WorkflowDataException if the answer is not a {@link ToolConfirmation} + */ + private @Nullable ToolConfirmation confirmationAnswer(Context context) { + if (confirmationPredicate == null) { + return null; + } + Object answer = context.resumeInputs().get(confirmationInterruptId(context)); + if (answer == null) { + return null; + } + if (!(answer instanceof Map map) + || !TOOL_CONFIRMATION_SCHEMA.properties().orElseThrow().keySet().containsAll(map.keySet()) + || (map.get("confirmed") != null && !(map.get("confirmed") instanceof Boolean)) + || (map.get("hint") != null && !(map.get("hint") instanceof String))) { + throw new WorkflowDataException( + "The answer to the confirmation request of node '" + + name() + + "' is not a ToolConfirmation."); + } + ToolConfirmation.Builder confirmation = ToolConfirmation.builder(); + if (map.get("hint") instanceof String hint) { + confirmation.hint(hint); + } + return confirmation + .confirmed(map.get("confirmed") instanceof Boolean confirmed && confirmed) + .payload(map.get("payload")) + .build(); + } + /** Returns the event or JSON value to send for {@code result}, or {@code null} for none. */ private @Nullable Object output(@Nullable O result) { String what = "output of node '" + name() + "'"; @@ -259,14 +392,16 @@ void cancel() { /** * Fluent builder for {@link FunctionNode}. Any property left unset falls back to the default of - * the matching {@link Node} property; name and body are required. + * the matching {@link Node} property, except that {@code rerunOnResume} defaults to true for a + * node that requires confirmation; name and body are required. */ public static final class Builder { private final PortCodec inputCodec; private final PortCodec outputCodec; private @Nullable String name; private String description = ""; - private boolean rerunOnResume; + private @Nullable Boolean rerunOnResume; + private @Nullable ConfirmationPredicate confirmationPredicate; private boolean waitForOutput; private NodeConfig config = NodeConfig.builder().build(); private @Nullable Schema inputSchema; @@ -293,7 +428,10 @@ public Builder description(String description) { return this; } - /** Sets whether a resume runs the node again rather than using the resuming answer. */ + /** + * Sets whether a resume runs the node again rather than using the resuming answer. Unset, it is + * true for a node that requires confirmation and false otherwise. + */ @CanIgnoreReturnValue public Builder rerunOnResume(boolean rerunOnResume) { this.rerunOnResume = rerunOnResume; @@ -340,6 +478,31 @@ public Builder stateSchema(@Nullable Schema stateSchema) { return this; } + /** + * Sets whether each run first asks the user to approve it; either overload replaces the other. + * Such a node reruns on resume, so a body that asks for input reads its answer from {@link + * Context#resumeInputs()} before asking. An approved run reads the user's decision from {@link + * Context#toolConfirmation()}, and a rejected run skips the body and outputs {@code {"error": + * "This tool call is rejected."}} whatever the output type, so a successor that must handle a + * rejection takes {@code Object}. + */ + @CanIgnoreReturnValue + public Builder requireConfirmation(boolean requireConfirmation) { + this.confirmationPredicate = requireConfirmation ? ALWAYS_CONFIRM : null; + return this; + } + + /** + * Sets a predicate that decides for each input whether a run asks the user to approve it, as + * the {@code boolean} overload does for every input. Once the node has asked, the answer + * decides. + */ + @CanIgnoreReturnValue + public Builder requireConfirmation(ConfirmationPredicate predicate) { + this.confirmationPredicate = checkNotNull(predicate); + return this; + } + /** Sets the function the node runs on each input. Required. */ @CanIgnoreReturnValue public Builder body(NodeBody body) { @@ -351,10 +514,19 @@ public Builder body(NodeBody body) { * Builds the node. * * @throws IllegalStateException if the name or the body is not set + * @throws WorkflowConfigurationException if the node requires confirmation but {@code + * rerunOnResume} is set to false */ public FunctionNode build() { checkState(name != null, "FunctionNode.Builder requires name to be set."); checkState(body != null, "FunctionNode.Builder for '%s' requires body to be set.", name); + if (confirmationPredicate != null && rerunOnResume != null && !rerunOnResume) { + throw new WorkflowConfigurationException( + "node '" + + name + + "' requires confirmation, so it needs rerunOnResume(true) to run again once the" + + " user answers."); + } return new FunctionNode<>(this); } } 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 5b6ef3044..2a8f0bebd 100644 --- a/core/src/main/java/com/google/adk/workflow/NodeExecutionState.java +++ b/core/src/main/java/com/google/adk/workflow/NodeExecutionState.java @@ -423,6 +423,11 @@ void adoptReplayedOutput(Object output) { node.name()); } + /** Returns whether a {@code useAsOutput} child holds the run's output, so its own is dropped. */ + boolean hasOutputDelegate() { + return outputState.get().status().isDelegated(); + } + /** * Returns the output a {@code useAsOutput} descendant gave this run, or that the run adopted on * replay, or empty if the run did not delegate its output or holds no output yet. 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 e88a7dbda..71d7ce201 100644 --- a/core/src/test/java/com/google/adk/workflow/FunctionNodeTest.java +++ b/core/src/test/java/com/google/adk/workflow/FunctionNodeTest.java @@ -22,6 +22,8 @@ import static com.google.adk.workflow.WorkflowResumeTesting.go; import static com.google.adk.workflow.WorkflowResumeTesting.longRunningAsk; import static com.google.adk.workflow.WorkflowResumeTesting.replyInput; +import static com.google.adk.workflow.WorkflowResumeTesting.requestAnswer; +import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.collect.ImmutableMap.toImmutableMap; import static com.google.common.collect.MoreCollectors.onlyElement; import static com.google.common.truth.Truth.assertThat; @@ -52,21 +54,28 @@ import com.fasterxml.jackson.databind.node.JsonNodeFactory; import com.fasterxml.jackson.databind.ser.std.ToStringSerializer; import com.fasterxml.jackson.databind.util.StdConverter; +import com.google.adk.agents.Context; import com.google.adk.apps.App; import com.google.adk.events.Event; import com.google.adk.events.EventActions; +import com.google.adk.events.ToolConfirmation; import com.google.adk.runner.Runner; import com.google.adk.sessions.Session; +import com.google.adk.testing.TestNode; +import com.google.adk.testing.TestUtils; import com.google.adk.workflow.WorkflowResumeTesting.WorkflowSession; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.util.concurrent.Uninterruptibles; import com.google.genai.types.Content; import com.google.genai.types.FinishReason; +import com.google.genai.types.FunctionCall; import com.google.genai.types.Part; import com.google.genai.types.Schema; import com.google.genai.types.Type; import com.google.genai.types.VideoMetadata; +import io.reactivex.rxjava3.core.Completable; +import io.reactivex.rxjava3.core.Flowable; import io.reactivex.rxjava3.core.Scheduler; import io.reactivex.rxjava3.disposables.Disposable; import io.reactivex.rxjava3.plugins.RxJavaPlugins; @@ -82,6 +91,7 @@ import java.time.Instant; import java.time.ZoneId; import java.util.ArrayList; +import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Locale; @@ -99,6 +109,7 @@ import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Predicate; import org.jspecify.annotations.Nullable; import org.junit.Test; import org.junit.runner.RunWith; @@ -110,6 +121,12 @@ 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 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"; + private static final ImmutableMap REJECTED = + ImmutableMap.of("error", "This tool call is rejected."); private static final Schema ANY = Schema.builder().build(); private static final Schema STRING = typed(Type.Known.STRING); private static final Schema INTEGER = typed(Type.Known.INTEGER); @@ -3546,6 +3563,812 @@ public void of_mapPortsWithAJsonValueKey_readTheKeyBack() { assertThat(result.outputAt("wf@1/read@1")).isEqualTo("t1"); } + @Test + public void builder_requireConfirmation_asksBeforeTheBody() { + AtomicInteger runs = new AtomicInteger(); + FunctionNode transfer = transfer(runs).requireConfirmation(true).build(); + + ImmutableList events = new WorkflowSession(chain(transfer)).turn(go()); + + FunctionCall request = confirmationRequest(events); + assertThat(request.id()).hasValue(TRANSFER_CONFIRMATION); + Map args = request.args().orElseThrow(); + assertThat(args) + .containsEntry( + RequestInput.MESSAGE_KEY, "Please approve or reject the tool call transfer()."); + assertThat((Map) args.get(RequestInput.PAYLOAD_KEY)) + .containsEntry("tool_name", "transfer"); + assertThat(runs.get()).isEqualTo(0); + } + + @Test + public void builder_requireConfirmation_approvedRunsTheBodyWithTheAnswer() { + AtomicReference<@Nullable ToolConfirmation> seen = new AtomicReference<>(); + FunctionNode transfer = + FunctionNodes.builder(Content.class, String.class) + .name("transfer") + .requireConfirmation(true) + .body((ctx, message) -> { + seen.set(ctx.toolConfirmation().orElse(null)); + return "sent"; + }) + .build(); + WorkflowSession session = new WorkflowSession(chain(transfer, captureInput("reply"))); + session.advance(go()); + Map answer = new HashMap<>(); + answer.put("confirmed", true); + answer.put("hint", null); + answer.put("payload", ImmutableMap.of("limit", 50)); + + ImmutableList events = session.turn(requestAnswer(TRANSFER_CONFIRMATION, answer)); + + assertThat(replyInput(events)).isEqualTo("sent"); + assertThat(seen.get().confirmed()).isTrue(); + assertThat(seen.get().hint()).isEmpty(); + assertThat(seen.get().payload()).isEqualTo(ImmutableMap.of("limit", 50)); + } + + @Test + public void builder_requireConfirmation_rejectedSkipsTheBodyAndOutputsTheRejection() { + AtomicInteger runs = new AtomicInteger(); + FunctionNode transfer = transfer(runs).requireConfirmation(true).build(); + WorkflowSession session = new WorkflowSession(chain(transfer, captureInput("reply"))); + session.advance(go()); + + ImmutableList events = session.turn(confirmation(TRANSFER_CONFIRMATION, false)); + + assertThat(replyInput(events)).isEqualTo(REJECTED); + assertThat(runs.get()).isEqualTo(0); + } + + @Test + public void resume_confirmationApprovedWithoutResumability_runsTheBodyOnceAndContinues() { + AtomicInteger runs = new AtomicInteger(); + FunctionNode transfer = transfer(runs).requireConfirmation(true).build(); + WorkflowSession session = + new WorkflowSession(chain(transfer, captureInput("reply")), /* resumable= */ false); + ImmutableList askEvents = session.turn(go()); + + ImmutableList events = session.turn(confirmation(TRANSFER_CONFIRMATION, true)); + + assertThat(confirmationRequest(askEvents).id()).hasValue(TRANSFER_CONFIRMATION); + assertThat(replyInput(events)).isEqualTo("sent"); + assertThat(runs.get()).isEqualTo(1); + } + + @Test + public void resume_confirmationRejectedWithoutResumability_outputsTheRejectionWithoutTheBody() { + AtomicInteger runs = new AtomicInteger(); + FunctionNode transfer = transfer(runs).requireConfirmation(true).build(); + WorkflowSession session = + new WorkflowSession(chain(transfer, captureInput("reply")), /* resumable= */ false); + ImmutableList askEvents = session.turn(go()); + + ImmutableList events = session.turn(confirmation(TRANSFER_CONFIRMATION, false)); + + assertThat(confirmationRequest(askEvents).id()).hasValue(TRANSFER_CONFIRMATION); + assertThat(replyInput(events)).isEqualTo(REJECTED); + assertThat(runs.get()).isEqualTo(0); + } + + @Test + public void resume_laterAskAfterARejection_replaysTheRejection() { + AtomicInteger runs = new AtomicInteger(); + FunctionNode transfer = transfer(runs).requireConfirmation(true).build(); + WorkflowSession session = new WorkflowSession(chain(transfer, new AskOnce("ask", "ask_1"))); + session.advance(go()); + session.advance(confirmation(TRANSFER_CONFIRMATION, false)); + + ImmutableList events = session.turn(requestAnswer("ask_1", ImmutableMap.of("ok", true))); + + assertThat(WorkflowResumeTesting.outputEventAt(events, "wf@1/ask@1").output()) + .hasValue(ImmutableMap.of("received", REJECTED)); + assertThat(confirmationRequests(events)).isEmpty(); + assertThat(runs.get()).isEqualTo(0); + } + + @Test + public void resume_laterAskAfterAnApprovedRunWithoutOutput_doesNotRunTheBodyAgain() { + AtomicInteger runs = new AtomicInteger(); + FunctionNode notify = + FunctionNodes.builder(Content.class, Void.class) + .name("notify") + .requireConfirmation(true) + .body((ctx, message) -> { + runs.incrementAndGet(); + return null; + }) + .build(); + Node ask = new AskOnce("ask", "ask_1"); + WorkflowSession session = + new WorkflowSession( + Workflow.builder() + .name("wf") + .edges(new Edge(Start.INSTANCE, notify), new Edge(Start.INSTANCE, ask)) + .build()); + session.advance(go()); + session.advance(confirmation("wf_tool_confirmation:wf@1/notify@1", true)); + + ImmutableList events = session.turn(requestAnswer("ask_1", ImmutableMap.of("ok", true))); + + assertThat(confirmationRequests(events)).isEmpty(); + assertThat(runs.get()).isEqualTo(1); + } + + @Test + public void resume_approvedBodyThatWaits_rerunsOnTheAnswerWithoutAskingAgain() { + AtomicInteger runs = new AtomicInteger(); + FunctionNode transfer = + FunctionNodes.builder(Content.class, Object.class) + .name("transfer") + .requireConfirmation(true) + .body((ctx, message) -> { + runs.incrementAndGet(); + Object reply = ctx.resumeInputs().get("ask_1"); + return reply != null ? reply : WorkflowResumeTesting.interrupt("ask_1"); + }) + .build(); + WorkflowSession session = new WorkflowSession(chain(transfer, captureInput("reply"))); + List events = new ArrayList<>(session.turn(go())); + events.addAll(session.turn(confirmation(TRANSFER_CONFIRMATION, true))); + + events.addAll(session.turn(answer("ask_1", ImmutableMap.of("ok", true)))); + + assertThat(confirmationRequests(events)).hasSize(1); + assertThat(replyInput(events)).isEqualTo(ImmutableMap.of("ok", true)); + assertThat(runs.get()).isEqualTo(2); + } + + @Test + public void resume_approvedBodyThatDelegatesItsOutput_writesNoEmptyEventAndIsNotRerun() { + AtomicInteger runs = new AtomicInteger(); + FunctionNode child = + FunctionNodes.of("child", Void.class, String.class, (ctx, unused) -> "child_done"); + FunctionNode parent = + FunctionNodes.builder(Content.class, String.class) + .name("parent") + .requireConfirmation(true) + .body((ctx, message) -> { + runs.incrementAndGet(); + Object unused = + ctx.runNode(child, null, RunNodeConfig.builder().useAsOutput(true).build()) + .blockingGet(); + return null; + }) + .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)); + + session.advance(requestAnswer("review", ImmutableMap.of("ok", true))); + + ImmutableList parentEvents = + confirmed.stream() + .filter(event -> WorkflowResumeTesting.pathOf(event).startsWith("wf@1/parent@1")) + .collect(toImmutableList()); + assertThat(WorkflowResumeTesting.pathsOf(parentEvents)) + .containsExactly("wf@1/parent@1/child@1"); + assertThat(parentEvents.get(0).output()).hasValue("child_done"); + assertThat(runs.get()).isEqualTo(1); + } + + @Test + public void resume_approvedBodyWhoseOnlyOutputIsDropped_isNotRerunWhenALaterNodeResumes() { + AtomicInteger runs = new AtomicInteger(); + AtomicInteger childRuns = new AtomicInteger(); + FunctionNode child = + FunctionNodes.of( + "child", + Void.class, + String.class, + (ctx, unused) -> { + childRuns.incrementAndGet(); + return null; + }); + FunctionNode parent = + FunctionNodes.builder(Content.class, String.class) + .name("parent") + .requireConfirmation(true) + .body((ctx, message) -> { + runs.incrementAndGet(); + Object unused = + ctx.runNode(child, null, RunNodeConfig.builder().useAsOutput(true).build()) + .blockingGet(); + return "fallback"; + }) + .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(requestAnswer("review", ImmutableMap.of("ok", true))); + + assertThat(runs.get()).isEqualTo(1); + assertThat(childRuns.get()).isEqualTo(1); + } + + @Test + public void resume_approvedBodyThatRoutesAfterAnEmptyDelegation_reachesTheRoutedNode() { + FunctionNode parent = + delegatingToASilentChild( + new AtomicInteger(), + Event.builder() + .actions( + EventActions.builder() + .route(ImmutableList.of(new Route.Tag("approved"))) + .build()) + .build()); + Node reply = captureInput("reply"); + WorkflowSession session = + new WorkflowSession( + Workflow.builder() + .name("wf") + .edges( + new Edge(Start.INSTANCE, parent), + new Edge(parent, reply, new Route.Tag("approved"))) + .build()); + session.advance(go()); + + ImmutableList events = + session.turn(confirmation("wf_tool_confirmation:wf@1/parent@1", true)); + + assertThat(WorkflowResumeTesting.pathsOf(events)).contains("wf@1/reply@1"); + } + + @Test + public void resume_approvedBodyThatSendsAMessageAfterAnEmptyDelegation_sendsItAndItsState() { + FunctionNode parent = + delegatingToASilentChild( + new AtomicInteger(), + Event.builder() + .content(Content.fromParts(Part.fromText("approved"))) + .actions( + EventActions.builder() + .stateDelta(new HashMap<>(ImmutableMap.of("my_key", "my_value"))) + .build()) + .build()); + FunctionNode read = + FunctionNodes.of( + "read", Object.class, String.class, (ctx, unused) -> "" + ctx.state().get("my_key")); + WorkflowSession session = new WorkflowSession(chain(parent, read)); + session.advance(go()); + + ImmutableList events = + session.turn(confirmation("wf_tool_confirmation:wf@1/parent@1", true)); + + Event sent = + events.stream() + .filter(event -> WorkflowResumeTesting.pathOf(event).equals("wf@1/parent@1")) + .collect(onlyElement()); + assertThat(sent.content().get().text()).isEqualTo("approved"); + assertThat(sent.output()).isEmpty(); + assertThat(WorkflowResumeTesting.outputEventAt(events, "wf@1/read@1").output()) + .hasValue("my_value"); + } + + @Test + public void resume_approvedBodyThatReturnsOnlyAnOutputAfterAnEmptyDelegation_isNotRerun() { + AtomicInteger runs = new AtomicInteger(); + FunctionNode parent = + 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(requestAnswer("review", ImmutableMap.of("ok", true))); + + assertThat(runs.get()).isEqualTo(1); + } + + @Test + public void resume_approvedBodyThatWaitsAfterAnEmptyDelegation_asksTheUser() { + FunctionNode parent = + delegatingToASilentChild(new AtomicInteger(), WorkflowResumeTesting.interrupt("ask_1")); + WorkflowSession session = new WorkflowSession(chain(parent, captureInput("reply"))); + session.advance(go()); + + ImmutableList events = + session.turn(confirmation("wf_tool_confirmation:wf@1/parent@1", true)); + + assertThat( + events.stream() + .flatMap(event -> event.functionCalls().stream()) + .map(call -> call.id().orElse(""))) + .contains("ask_1"); + assertThat(WorkflowResumeTesting.pathsOf(events)).doesNotContain("wf@1/reply@1"); + } + + @Test + public void builder_requireConfirmationPredicateTrue_asks() { + FunctionNode pay = pay(amount -> amount > 100); + + ImmutableList events = + new WorkflowSession(chain(emitter("amount", 500), pay)).turn(go()); + + FunctionCall request = confirmationRequest(events); + assertThat(request.id()).hasValue(PAY_CONFIRMATION); + Map args = request.args().orElseThrow(); + assertThat((Map) args.get(RequestInput.PAYLOAD_KEY)).containsEntry("args", 500); + assertThat(args.get(RequestInput.RESPONSE_SCHEMA_KEY)) + .isEqualTo( + ImmutableMap.of( + "type", + "OBJECT", + "properties", + ImmutableMap.of( + "hint", ImmutableMap.of("type", "STRING", "nullable", true), + "confirmed", ImmutableMap.of("type", "BOOLEAN", "nullable", true), + "payload", ImmutableMap.of()))); + } + + @Test + public void builder_requireConfirmationPredicateFalse_runsTheBodyWithoutAsking() { + FunctionNode pay = pay(amount -> amount > 100); + + ImmutableList events = new WorkflowSession(chain(emitter("amount", 5), pay)).turn(go()); + + assertThat(confirmationRequests(events)).isEmpty(); + assertThat(WorkflowResumeTesting.outputEventAt(events, "wf@1/pay@1").output()) + .hasValue("paid 5"); + } + + @Test + public void builder_requireConfirmationPredicate_runsOncePerAsk() { + AtomicInteger calls = new AtomicInteger(); + FunctionNode pay = pay(amount -> calls.incrementAndGet() > 0); + WorkflowSession session = new WorkflowSession(chain(emitter("amount", 500), pay)); + session.advance(go()); + + ImmutableList events = session.turn(confirmation(PAY_CONFIRMATION, true)); + + assertThat(WorkflowResumeTesting.outputEventAt(events, "wf@1/pay@1").output()) + .hasValue("paid 500"); + assertThat(calls.get()).isEqualTo(1); + } + + @Test + public void builder_requireConfirmationPredicateThrows_failsTheRunWithoutTheBody() { + AtomicInteger runs = new AtomicInteger(); + FunctionNode transfer = + transfer(runs) + .requireConfirmation((ctx, message) -> { + throw new IOException("policy unavailable"); + }) + .build(); + + RunResult result = run(chain(transfer)); + + assertThat(result.errorMessage("IOException")).contains("policy unavailable"); + assertThat(runs.get()).isEqualTo(0); + } + + @Test + public void resume_answerThatIsNotAToolConfirmation_failsTheResume() { + FunctionNode transfer = + transfer(new AtomicInteger()).requireConfirmation(true).build(); + WorkflowSession session = new WorkflowSession(chain(transfer)); + session.advance(go()); + Content answer = requestAnswer(TRANSFER_CONFIRMATION, ImmutableMap.of("confirmed", "yes")); + + IllegalArgumentException error = + assertThrows(IllegalArgumentException.class, () -> session.turn(answer)); + + assertThat(error.getClass()).isEqualTo(IllegalArgumentException.class); + assertThat(error) + .hasMessageThat() + .startsWith("Validation failed for interrupt " + TRANSFER_CONFIRMATION + ":"); + } + + @Test + public void builder_requireConfirmation_defaultsRerunOnResumeToTrue() { + FunctionNode transfer = + transfer(new AtomicInteger()).requireConfirmation(true).build(); + + assertThat(transfer.rerunOnResume()).isTrue(); + } + + @Test + public void builder_requireConfirmationWithRerunOnResumeFalse_failsAtBuild() { + FunctionNode.Builder transfer = + transfer(new AtomicInteger()).rerunOnResume(false).requireConfirmation(true); + + WorkflowConfigurationException error = + assertThrows(WorkflowConfigurationException.class, transfer::build); + + assertThat(error).hasMessageThat().contains("rerunOnResume(true)"); + } + + @Test + public void builder_requireConfirmationFalse_clearsAnEarlierPredicate() { + FunctionNode pay = + pay(amount -> true).toBuilder().requireConfirmation(false).build(); + + ImmutableList events = + new WorkflowSession(chain(emitter("amount", 500), pay)).turn(go()); + + assertThat(confirmationRequests(events)).isEmpty(); + } + + @Test + public void builder_predicateAfterRequireConfirmationTrue_predicateDecides() { + AtomicInteger runs = new AtomicInteger(); + FunctionNode transfer = + transfer(runs).requireConfirmation(true).requireConfirmation((ctx, m) -> false).build(); + + ImmutableList events = new WorkflowSession(chain(transfer)).turn(go()); + + assertThat(confirmationRequests(events)).isEmpty(); + assertThat(runs.get()).isEqualTo(1); + } + + @Test + public void resume_confirmedNull_rejectsTheRun() { + AtomicInteger runs = new AtomicInteger(); + FunctionNode transfer = transfer(runs).requireConfirmation(true).build(); + WorkflowSession session = new WorkflowSession(chain(transfer, captureInput("reply"))); + session.advance(go()); + Map answer = new HashMap<>(); + answer.put("confirmed", null); + + ImmutableList events = session.turn(requestAnswer(TRANSFER_CONFIRMATION, answer)); + + assertThat(replyInput(events)).isEqualTo(REJECTED); + assertThat(runs.get()).isEqualTo(0); + } + + @Test + public void resume_answerWithAHint_bodySeesTheHint() { + AtomicReference<@Nullable ToolConfirmation> seen = new AtomicReference<>(); + FunctionNode transfer = + FunctionNodes.builder(Content.class, String.class) + .name("transfer") + .requireConfirmation(true) + .body((ctx, message) -> { + seen.set(ctx.toolConfirmation().orElse(null)); + return "sent"; + }) + .build(); + WorkflowSession session = new WorkflowSession(chain(transfer)); + session.advance(go()); + + session.advance( + requestAnswer( + TRANSFER_CONFIRMATION, ImmutableMap.of("confirmed", true, "hint", "under the limit"))); + + assertThat(seen.get().hint()).isEqualTo("under the limit"); + } + + @Test + public void runNode_answerThatIsNotAMap_throwsWorkflowDataException() { + FunctionNode transfer = + transfer(new AtomicInteger()).requireConfirmation(true).build(); + Context context = nodeContext(ImmutableMap.of("wf_tool_confirmation:transfer@1", "approved")); + Flowable run = transfer.runNode(context, null); + + WorkflowDataException error = assertThrows(WorkflowDataException.class, run::blockingFirst); + + assertThat(error) + .hasMessageThat() + .isEqualTo( + "The answer to the confirmation request of node 'transfer' is not a ToolConfirmation."); + } + + @Test + public void runNode_answerWithConfirmedThatIsNotABoolean_throwsWorkflowDataException() { + FunctionNode transfer = + transfer(new AtomicInteger()).requireConfirmation(true).build(); + Context context = + nodeContext( + ImmutableMap.of( + "wf_tool_confirmation:transfer@1", ImmutableMap.of("confirmed", "yes"))); + Flowable run = transfer.runNode(context, null); + + WorkflowDataException error = assertThrows(WorkflowDataException.class, run::blockingFirst); + + assertThat(error) + .hasMessageThat() + .isEqualTo( + "The answer to the confirmation request of node 'transfer' is not a ToolConfirmation."); + } + + @Test + public void runNode_answerWithHintThatIsNotAString_throwsWorkflowDataException() { + FunctionNode transfer = + transfer(new AtomicInteger()).requireConfirmation(true).build(); + Context context = + nodeContext( + ImmutableMap.of( + "wf_tool_confirmation:transfer@1", ImmutableMap.of("confirmed", true, "hint", 5))); + Flowable run = transfer.runNode(context, null); + + WorkflowDataException error = assertThrows(WorkflowDataException.class, run::blockingFirst); + + assertThat(error) + .hasMessageThat() + .isEqualTo( + "The answer to the confirmation request of node 'transfer' is not a ToolConfirmation."); + } + + @Test + public void runNode_answerWithAKeyOutsideTheSchema_throwsWorkflowDataException() { + FunctionNode transfer = + transfer(new AtomicInteger()).requireConfirmation(true).build(); + Context context = + nodeContext( + ImmutableMap.of( + "wf_tool_confirmation:transfer@1", + ImmutableMap.of("confirmed", true, "approved", true))); + Flowable run = transfer.runNode(context, null); + + WorkflowDataException error = assertThrows(WorkflowDataException.class, run::blockingFirst); + + assertThat(error) + .hasMessageThat() + .isEqualTo( + "The answer to the confirmation request of node 'transfer' is not a ToolConfirmation."); + } + + @Test + public void toBuilder_requireConfirmationOnANodeFromOf_asks() { + FunctionNode echo = + FunctionNodes.of("transfer", Content.class, String.class, (ctx, message) -> "sent"); + FunctionNode confirming = echo.toBuilder().requireConfirmation(true).build(); + + ImmutableList events = new WorkflowSession(chain(confirming)).turn(go()); + + assertThat(confirmationRequest(events).id()).hasValue(TRANSFER_CONFIRMATION); + } + + @Test + public void toBuilder_nodeWithAConfirmationPredicate_copyAsks() { + FunctionNode copy = pay(amount -> amount > 100).toBuilder().build(); + + ImmutableList events = + new WorkflowSession(chain(emitter("amount", 500), copy)).turn(go()); + + assertThat(confirmationRequest(events).id()).hasValue(PAY_CONFIRMATION); + } + + @Test + public void toBuilder_rerunOnResume_copiesOnlyAnExplicitValue() { + FunctionNode derived = + transfer(new AtomicInteger()).requireConfirmation(true).build(); + FunctionNode explicitTrue = + transfer(new AtomicInteger()).rerunOnResume(true).requireConfirmation(true).build(); + FunctionNode explicitFalse = + transfer(new AtomicInteger()).rerunOnResume(false).build(); + + FunctionNode derivedCopy = + derived.toBuilder().requireConfirmation(false).build(); + FunctionNode explicitTrueCopy = + explicitTrue.toBuilder().requireConfirmation(false).build(); + FunctionNode.Builder explicitFalseCopy = + explicitFalse.toBuilder().requireConfirmation(true); + + assertThat(derivedCopy.rerunOnResume()).isFalse(); + assertThat(explicitTrueCopy.rerunOnResume()).isTrue(); + assertThrows(WorkflowConfigurationException.class, explicitFalseCopy::build); + } + + @Test + public void resume_rejectionIntoAStringPort_failsTheSuccessor() { + FunctionNode transfer = + transfer(new AtomicInteger()).requireConfirmation(true).build(); + FunctionNode read = + FunctionNodes.of("read", String.class, String.class, (ctx, text) -> text); + WorkflowSession session = new WorkflowSession(chain(transfer, read)); + session.advance(go()); + Content rejection = confirmation(TRANSFER_CONFIRMATION, false); + + NodeInputValidationException error = + assertThrows(NodeInputValidationException.class, () -> session.turn(rejection)); + + assertThat(error).hasMessageThat().contains("input of node 'read' does not match"); + } + + @Test + public void resume_rejectionIntoARecordWithRequiredComponents_failsTheSuccessor() { + FunctionNode order = + FunctionNodes.builder(Content.class, Order.class) + .name("transfer") + .requireConfirmation(true) + .body((ctx, message) -> new Order("o1", 1)) + .build(); + FunctionNode describe = + FunctionNodes.of("describe", Order.class, String.class, (ctx, o) -> o.id()); + WorkflowSession session = new WorkflowSession(chain(order, describe)); + session.advance(go()); + Content rejection = confirmation(TRANSFER_CONFIRMATION, false); + + NodeInputValidationException error = + assertThrows(NodeInputValidationException.class, () -> session.turn(rejection)); + + assertThat(error).hasMessageThat().contains("input of node 'describe' does not match"); + } + + @Test + public void resume_rejectionIntoARecordOfOptionalComponents_leavesThemUnset() { + FunctionNode note = + FunctionNodes.builder(Content.class, OptionalNote.class) + .name("transfer") + .requireConfirmation(true) + .body((ctx, message) -> new OptionalNote("ok")) + .build(); + AtomicReference<@Nullable OptionalNote> seen = new AtomicReference<>(); + FunctionNode read = + FunctionNodes.of( + "read", + OptionalNote.class, + String.class, + (ctx, r) -> { + seen.set(r); + return "read"; + }); + WorkflowSession session = new WorkflowSession(chain(note, read)); + session.advance(go()); + + session.advance(confirmation(TRANSFER_CONFIRMATION, false)); + + assertThat(seen.get()).isEqualTo(new OptionalNote(null)); + } + + @Test + public void parallelWorkers_confirmingItems_askEachAndRunOnlyTheApprovedOnes() { + List runs = Collections.synchronizedList(new ArrayList<>()); + FunctionNode pay = + FunctionNodes.builder(Integer.class, String.class) + .name("pay") + .requireConfirmation(true) + .body((ctx, amount) -> { + runs.add(amount); + return "paid " + amount; + }) + .build(); + WorkflowSession session = + new WorkflowSession( + chain( + emitter("amounts", ImmutableList.of(1, 2)), + ParallelWorkers.parallelNode(pay), + captureInput("reply"))); + ImmutableList asked = session.turn(go()); + session.advance(confirmation("wf_tool_confirmation:wf@1/pay@1/pay@1", true)); + + ImmutableList events = + session.turn(confirmation("wf_tool_confirmation:wf@1/pay@1/pay@2", false)); + + assertThat( + confirmationRequests(asked).stream() + .map(call -> call.id().orElseThrow()) + .collect(toImmutableList())) + .containsExactly( + "wf_tool_confirmation:wf@1/pay@1/pay@1", "wf_tool_confirmation:wf@1/pay@1/pay@2"); + assertThat(replyInput(events)).isEqualTo(ImmutableList.of("paid 1", REJECTED)); + assertThat(runs).containsExactly(1); + } + + @Test + public void runNodeChild_confirmationApproved_runsTheChildOnceAndTheParentOutputsIt() { + AtomicInteger runs = new AtomicInteger(); + FunctionNode transfer = transfer(runs).requireConfirmation(true).build(); + WorkflowSession session = new WorkflowSession(chain(parent(transfer), captureInput("reply"))); + ImmutableList asked = session.turn(go()); + + ImmutableList events = session.turn(confirmation(CHILD_TRANSFER_CONFIRMATION, true)); + + assertThat(confirmationRequest(asked).id()).hasValue(CHILD_TRANSFER_CONFIRMATION); + assertThat(replyInput(events)).isEqualTo("sent"); + assertThat(runs.get()).isEqualTo(1); + } + + @Test + public void runNodeChild_confirmationRejected_parentOutputsTheRejectionWithoutTheChildBody() { + AtomicInteger runs = new AtomicInteger(); + FunctionNode transfer = transfer(runs).requireConfirmation(true).build(); + WorkflowSession session = new WorkflowSession(chain(parent(transfer), captureInput("reply"))); + session.advance(go()); + + ImmutableList events = session.turn(confirmation(CHILD_TRANSFER_CONFIRMATION, false)); + + assertThat(replyInput(events)).isEqualTo(REJECTED); + assertThat(runs.get()).isEqualTo(0); + } + + record OptionalNote(@Nullable String note) {} + + /** Returns a builder for a node {@code transfer} that counts its runs and outputs "sent". */ + private static FunctionNode.Builder transfer(AtomicInteger runs) { + return FunctionNodes.builder(Content.class, String.class) + .name("transfer") + .body((ctx, message) -> { + runs.incrementAndGet(); + return "sent"; + }); + } + + /** Returns a node {@code parent} that runs {@code child} on its input and outputs its output. */ + private static FunctionNode parent(Node child) { + return FunctionNodes.builder(Content.class, Object.class) + .name("parent") + .rerunOnResume(true) + .body((ctx, message) -> ctx.runNode(child, message).blockingGet()) + .build(); + } + + /** + * Returns a node {@code parent} that requires confirmation and, once approved, gives its output + * to a child that outputs nothing, then returns {@code event}; {@code runs} counts its body's + * runs. + */ + private static FunctionNode delegatingToASilentChild( + AtomicInteger runs, Event event) { + FunctionNode child = + FunctionNodes.of("child", Void.class, String.class, (ctx, unused) -> null); + return FunctionNodes.builder(Content.class, Event.class) + .name("parent") + .requireConfirmation(true) + .body((ctx, message) -> { + runs.incrementAndGet(); + Object unused = + ctx.runNode(child, null, RunNodeConfig.builder().useAsOutput(true).build()) + .blockingGet(); + return event; + }) + .build(); + } + + /** Returns the context of a run of node {@code transfer} that resumes with {@code answers}. */ + private static Context nodeContext(Map answers) { + return new Context( + TestUtils.createInvocationContext(TestUtils.createRootAgent()), + NodeExecutionState.builder() + .node(TestNode.of("transfer")) + .eventSink((event, nodeFailure) -> Completable.complete()) + .runId("1") + .resumeInputs(answers) + .build()); + } + + /** Returns a node {@code pay} that asks for confirmation when {@code asks} accepts the amount. */ + private static FunctionNode pay(Predicate asks) { + return FunctionNodes.builder(Integer.class, String.class) + .name("pay") + .requireConfirmation((ctx, amount) -> asks.test(amount)) + .body((ctx, amount) -> "paid " + amount) + .build(); + } + + private static Content confirmation(String interruptId, boolean confirmed) { + return requestAnswer(interruptId, ImmutableMap.of("confirmed", confirmed)); + } + + private static ImmutableList confirmationRequests(List events) { + return events.stream() + .flatMap(event -> event.functionCalls().stream()) + .filter(call -> call.id().orElse("").startsWith("wf_tool_confirmation:")) + .collect(toImmutableList()); + } + + private static FunctionCall confirmationRequest(List events) { + return confirmationRequests(events).stream().collect(onlyElement()); + } + + /** A node that asks for input {@code interruptId}, then outputs its input under "received". */ + private record AskOnce(String name, String interruptId) implements Node { + + @Override + public boolean rerunOnResume() { + return true; + } + + @Override + public Flowable runNode(Context context, @Nullable Object nodeInput) { + return context.resumeInputs().containsKey(interruptId) + ? Flowable.just(Collections.singletonMap("received", nodeInput)) + : Flowable.just(RequestInput.builder().interruptId(interruptId).build()); + } + } + private static Optional inputSchema(Class type) { return FunctionNodes.of("n", type, Void.class, (ctx, input) -> null).inputSchema(); }