From b7baf089c6ae0ee158ffd8772e42e3074e3b972c Mon Sep 17 00:00:00 2001 From: Mateusz Krawiec Date: Fri, 13 Nov 2026 10:38:36 -0800 Subject: [PATCH] fix(workflow): treat a null confirmed in a tool node's answer as a rejection PiperOrigin-RevId: 1021150225 --- .../com/google/adk/workflow/BaseNode.java | 36 +++-- .../adk/workflow/ConfirmationPredicate.java | 3 +- .../com/google/adk/workflow/FunctionNode.java | 141 ++++++------------ .../adk/workflow/ToolConfirmations.java | 123 +++++++++++++++ .../com/google/adk/workflow/ToolNodes.java | 114 ++------------ .../google/adk/workflow/FunctionNodeTest.java | 137 +++++++++++++++-- .../google/adk/workflow/ToolNodesTest.java | 42 +++++- 7 files changed, 365 insertions(+), 231 deletions(-) create mode 100644 core/src/main/java/com/google/adk/workflow/ToolConfirmations.java diff --git a/core/src/main/java/com/google/adk/workflow/BaseNode.java b/core/src/main/java/com/google/adk/workflow/BaseNode.java index ce533cbcf..4a7a3bc7d 100644 --- a/core/src/main/java/com/google/adk/workflow/BaseNode.java +++ b/core/src/main/java/com/google/adk/workflow/BaseNode.java @@ -132,22 +132,26 @@ protected void validateOutgoingEdges(Collection edges) {} */ final Flowable run(Context context, @Nullable Object nodeInput) { InvocationContext invocationContext = checkNotNull(context).invocationContext(); - // The node runner fills in the author, as only it knows the node's place in the graph. return Flowable.defer(() -> runNode(context, nodeInput)) - .map( - item -> { - if (item instanceof Event event) { - return event; - } - if (item instanceof RequestInput request) { - return request.toEvent(invocationContext); - } - return Event.builder() - .id(invocationContext.newUuid()) - .timestamp(invocationContext.now().toEpochMilli()) - .author("") - .output(item) - .build(); - }); + .map(item -> { + if (item instanceof Event event) { + return event; + } + if (item instanceof RequestInput request) { + return request.toEvent(invocationContext); + } + return newEventBuilder(invocationContext).output(item).build(); + }); + } + + /** + * Returns a builder of an event with a new ID and the current time. The node runner fills in the + * author, as only it knows the node's place in the graph. + */ + static Event.Builder newEventBuilder(InvocationContext invocationContext) { + return Event.builder() + .id(invocationContext.newUuid()) + .timestamp(invocationContext.now().toEpochMilli()) + .author(""); } } 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..88969fba1 100644 --- a/core/src/main/java/com/google/adk/workflow/FunctionNode.java +++ b/core/src/main/java/com/google/adk/workflow/FunctionNode.java @@ -28,12 +28,9 @@ 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; @@ -56,25 +53,12 @@ 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; @@ -171,17 +155,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,34 +185,51 @@ 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); + String what = "input of node '" + name() + "'"; + I input = inputCodec.decode(nodeInput, what); + ToolConfirmation confirmation = + confirmationPredicate != null ? ToolConfirmations.answer(context, name()) : null; 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)) { + RequestInput request = + ToolConfirmations.request( + context, + name(), + inputCodec.encode(input, what), + /* hint= */ null, + /* confirmationPayload= */ null); + return new RunOutcome<>(null, request, 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() + : newEventBuilder(context.invocationContext()).build(); } - if (output == null || state.hasOutputDelegate()) { - return finishedEvent(context); + if (output == null) { + return newEventBuilder(context.invocationContext()).build(); } } // After a delegation, NodeRunner drops a returned output. @@ -236,7 +240,7 @@ private void runBody(Context context, @Nullable Object nodeInput, MaybeEmitter( + private record RunOutcome( @Nullable ToolConfirmation confirmation, @Nullable Object replacement, @Nullable O result) {} /** @@ -273,65 +277,6 @@ 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() + "'"; diff --git a/core/src/main/java/com/google/adk/workflow/ToolConfirmations.java b/core/src/main/java/com/google/adk/workflow/ToolConfirmations.java new file mode 100644 index 000000000..e04eeb92f --- /dev/null +++ b/core/src/main/java/com/google/adk/workflow/ToolConfirmations.java @@ -0,0 +1,123 @@ +/* + * 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 static com.google.common.base.Strings.isNullOrEmpty; + +import com.google.adk.agents.Context; +import com.google.adk.events.ToolConfirmation; +import com.google.common.collect.ImmutableMap; +import com.google.genai.types.Schema; +import com.google.genai.types.Type; +import java.util.LinkedHashMap; +import java.util.Map; +import org.jspecify.annotations.Nullable; + +/** + * Static utilities for the confirmation request that a function or tool node run sends to ask the + * user to approve it, and for the user's answer. + */ +final class ToolConfirmations { + + /** The interrupt ID prefix of a confirmation request, followed by the node path. */ + private static final String INTERRUPT_ID_PREFIX = "wf_tool_confirmation:"; + + /** The schema of an answer to a confirmation request. */ + private static final Schema RESPONSE_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(); + + /** Returns the interrupt ID of the confirmation request of the node run in {@code context}. */ + private static String interruptId(Context context) { + return INTERRUPT_ID_PREFIX + context.requireNodeState().nodePath(); + } + + /** + * Returns the request that asks the user to approve calling {@code toolName} with {@code args}. A + * non-empty {@code hint} replaces the default message, and a non-null {@code confirmationPayload} + * goes into the request's payload. + */ + static RequestInput request( + Context context, + String toolName, + @Nullable Object args, + @Nullable String hint, + @Nullable Object confirmationPayload) { + Map payload = new LinkedHashMap<>(); + payload.put("tool_name", toolName); + payload.put("args", args); + if (confirmationPayload != null) { + payload.put("confirmation_payload", confirmationPayload); + } + return RequestInput.builder() + .interruptId(interruptId(context)) + .message( + isNullOrEmpty(hint) + ? "Please approve or reject the tool call " + toolName + "()." + : hint) + .payload(payload) + .responseSchema(RESPONSE_SCHEMA) + .build(); + } + + /** + * Returns the user's answer to the confirmation request of the node run in {@code context}, or + * null if there is none yet. A {@code confirmed} left out or null rejects the call, and a {@code + * hint} left out or null reads as an empty hint. + * + * @throws WorkflowDataException if the answer is not a {@link ToolConfirmation} + */ + static @Nullable ToolConfirmation answer(Context context, String nodeName) { + Object answer = context.resumeInputs().get(interruptId(context)); + if (answer == null) { + return null; + } + if (!(answer instanceof Map map) + || !RESPONSE_SCHEMA.properties().orElseThrow().keySet().containsAll(map.keySet())) { + throw notAToolConfirmation(nodeName); + } + Object confirmed = map.get("confirmed"); + Object hint = map.get("hint"); + if ((confirmed != null && !(confirmed instanceof Boolean)) + || (hint != null && !(hint instanceof String))) { + throw notAToolConfirmation(nodeName); + } + ToolConfirmation.Builder confirmation = ToolConfirmation.builder(); + if (hint instanceof String text) { + confirmation.hint(text); + } + return confirmation + .confirmed(Boolean.TRUE.equals(confirmed)) + .payload(map.get("payload")) + .build(); + } + + private static WorkflowDataException notAToolConfirmation(String nodeName) { + return new WorkflowDataException( + "The answer to the confirmation request of node '" + + nodeName + + "' is not a ToolConfirmation."); + } + + private ToolConfirmations() {} +} diff --git a/core/src/main/java/com/google/adk/workflow/ToolNodes.java b/core/src/main/java/com/google/adk/workflow/ToolNodes.java index 439958a6e..a473e4721 100644 --- a/core/src/main/java/com/google/adk/workflow/ToolNodes.java +++ b/core/src/main/java/com/google/adk/workflow/ToolNodes.java @@ -16,7 +16,6 @@ package com.google.adk.workflow; -import static com.google.common.base.Strings.isNullOrEmpty; import static com.google.common.base.Strings.nullToEmpty; import com.fasterxml.jackson.core.JsonProcessingException; @@ -26,9 +25,7 @@ import com.fasterxml.jackson.databind.ObjectReader; import com.google.adk.JsonBaseModel; import com.google.adk.agents.Context; -import com.google.adk.agents.InvocationContext; import com.google.adk.annotations.Experimental; -import com.google.adk.events.Event; import com.google.adk.events.EventActions; import com.google.adk.events.ToolConfirmation; import com.google.adk.plugins.Plugin; @@ -36,7 +33,6 @@ import com.google.adk.tools.ToolContext; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; -import com.google.common.collect.ImmutableSet; import com.google.errorprone.annotations.CanIgnoreReturnValue; import com.google.genai.types.Content; import com.google.genai.types.FunctionDeclaration; @@ -95,33 +91,6 @@ private static final class ToolNode extends BaseNode { .readerFor(Object.class) .with(DeserializationFeature.FAIL_ON_TRAILING_TOKENS); - /** The interrupt ID prefix of a run's confirmation request, followed by the node path. */ - private static final String CONFIRMATION_INTERRUPT_PREFIX = "wf_tool_confirmation:"; - - private static final String TOOL_NAME_KEY = "tool_name"; - private static final String ARGS_KEY = "args"; - private static final String CONFIRMATION_PAYLOAD_KEY = "confirmation_payload"; - private static final String CONFIRMATION_MESSAGE_FORMAT = - "Please approve or reject the tool call %s()."; - - private static final String ANSWER_CONFIRMED_KEY = "confirmed"; - private static final String ANSWER_HINT_KEY = "hint"; - private static final String ANSWER_PAYLOAD_KEY = "payload"; - private static final ImmutableSet ANSWER_KEYS = - ImmutableSet.of(ANSWER_CONFIRMED_KEY, ANSWER_HINT_KEY, ANSWER_PAYLOAD_KEY); - - /** The schema of the user's answer to a confirmation request. */ - private static final Schema TOOL_CONFIRMATION_SCHEMA = - Schema.builder() - .type(Type.Known.OBJECT) - .properties( - ImmutableMap.of( - ANSWER_CONFIRMED_KEY, Schema.builder().type(Type.Known.BOOLEAN).build(), - ANSWER_HINT_KEY, - Schema.builder().type(Type.Known.STRING).nullable(true).build(), - ANSWER_PAYLOAD_KEY, Schema.builder().build())) - .build(); - private final BaseTool tool; ToolNode(BaseTool tool, String name, NodeConfig config) { @@ -141,7 +110,11 @@ private static final class ToolNode extends BaseNode { @Override public Flowable runNode(Context context, @Nullable Object nodeInput) { return Single.defer(() -> runTool(context, nodeInput)) - .map(result -> returnedNothing(result) ? eventWithoutOutput(context) : result) + // A void tool's result becomes an event without output, so a resume does not rerun it. + .map(result -> + returnedNothing(result) + ? newEventBuilder(context.invocationContext()).build() + : result) .toFlowable(); } @@ -150,19 +123,6 @@ private boolean returnedNothing(Object result) { return result instanceof Map map && map.isEmpty() && declaresNoResult(tool); } - /** - * Returns an event without output for a tool declared to return nothing, whose result is still - * an empty map. The event records the run, so a resumed workflow does not run the tool again. - */ - private static Event eventWithoutOutput(Context context) { - InvocationContext invocationContext = context.invocationContext(); - return Event.builder() - .id(invocationContext.newUuid()) - .timestamp(invocationContext.now().toEpochMilli()) - .author("") - .build(); - } - /** Returns the tool's result, or the {@link RequestInput} that asks to confirm the call. */ private Single runTool(Context context, @Nullable Object nodeInput) { Parameters parameters = Parameters.of(tool); @@ -172,13 +132,12 @@ private Single runTool(Context context, @Nullable Object nodeInput) { args.put(param, context.state().get(param)); } } - String interruptId = CONFIRMATION_INTERRUPT_PREFIX + context.requireNodeState().nodePath(); ToolContext toolContext = ToolContext.builder(context.invocationContext()) .functionCallId(context.invocationContext().newUuid()) - .toolConfirmation(confirmationAnswer(context.resumeInputs().get(interruptId))) + .toolConfirmation(ToolConfirmations.answer(context, name())) .build(); - return callTool(context.invocationContext().pluginManager(), args, toolContext, interruptId) + return callTool(context, args, toolContext) // Kept on failure too, so the recorded deltas match the live session state. .doOnEvent((unusedResult, unusedError) -> keepDeltas(context, toolContext)); } @@ -189,14 +148,17 @@ private Single runTool(Context context, @Nullable Object nodeInput) { * callbacks. */ private Single callTool( - Plugin plugins, Map args, ToolContext toolContext, String interruptId) { + Context context, Map args, ToolContext toolContext) { + Plugin plugins = context.invocationContext().pluginManager(); Single toolCall = Single.defer(() -> tool.runAsync(args, toolContext)) .flatMap( result -> { ToolConfirmation requested = takeRequestedConfirmation(toolContext); return requested != null - ? Single.just(confirmationRequest(args, requested, interruptId)) + ? Single.just( + ToolConfirmations.request( + context, tool.name(), args, requested.hint(), requested.payload())) : afterToolCallbacks(plugins, args, toolContext, result); }, error -> { @@ -224,34 +186,6 @@ private Single afterToolCallbacks( .cast(Object.class); } - /** - * Returns the user's answer to this run's confirmation request, or null if there is none yet. - * An answer that leaves out {@code confirmed} rejects the call. - * - * @throws WorkflowDataException if the answer is not a {@link ToolConfirmation} - */ - private @Nullable ToolConfirmation confirmationAnswer(@Nullable Object answer) { - if (answer == null) { - return null; - } - if (!(answer instanceof Map map) - || !ANSWER_KEYS.containsAll(map.keySet()) - || (map.get(ANSWER_CONFIRMED_KEY) != null - && !(map.get(ANSWER_CONFIRMED_KEY) instanceof Boolean)) - || (map.get(ANSWER_HINT_KEY) != null && !(map.get(ANSWER_HINT_KEY) instanceof String))) { - throw new WorkflowDataException( - String.format( - "The answer to the confirmation request of tool node '%s' is not a" - + " ToolConfirmation.", - name())); - } - return ToolConfirmation.builder() - .confirmed(Boolean.TRUE.equals(map.get(ANSWER_CONFIRMED_KEY))) - .hint((String) map.get(ANSWER_HINT_KEY)) - .payload(map.get(ANSWER_PAYLOAD_KEY)) - .build(); - } - /** Removes and returns the confirmation the call asked for, or null if it asked for none. */ @CanIgnoreReturnValue private static @Nullable ToolConfirmation takeRequestedConfirmation(ToolContext toolContext) { @@ -261,30 +195,6 @@ private Single afterToolCallbacks( .remove(toolContext.functionCallId().orElseThrow()); } - /** - * Returns the request that asks the user to approve or reject this call, with the hint and - * payload in {@code requested}. - */ - private RequestInput confirmationRequest( - Map args, ToolConfirmation requested, String interruptId) { - ImmutableMap.Builder payload = - ImmutableMap.builder() - .put(TOOL_NAME_KEY, tool.name()) - .put(ARGS_KEY, args); - if (requested.payload() != null) { - payload.put(CONFIRMATION_PAYLOAD_KEY, requested.payload()); - } - return RequestInput.builder() - .interruptId(interruptId) - .message( - isNullOrEmpty(requested.hint()) - ? String.format(CONFIRMATION_MESSAGE_FORMAT, tool.name()) - : requested.hint()) - .payload(payload.buildOrThrow()) - .responseSchema(TOOL_CONFIRMATION_SCHEMA) - .build(); - } - private Map argumentsFrom(@Nullable Object nodeInput, Parameters parameters) { Object input = nodeInput instanceof Content content ? content.text() : nodeInput; if (input instanceof String text) { 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 diff --git a/core/src/test/java/com/google/adk/workflow/ToolNodesTest.java b/core/src/test/java/com/google/adk/workflow/ToolNodesTest.java index 309d4d9eb..6531eca5c 100644 --- a/core/src/test/java/com/google/adk/workflow/ToolNodesTest.java +++ b/core/src/test/java/com/google/adk/workflow/ToolNodesTest.java @@ -74,6 +74,7 @@ import java.util.Map; import java.util.Optional; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; import java.util.function.BiFunction; import org.jspecify.annotations.Nullable; import org.junit.Test; @@ -398,7 +399,7 @@ public void runNode_answerThatIsNotAConfirmation_failsTheNode() { WorkflowDataException thrown = assertThrows(WorkflowDataException.class, () -> run(tool, context, "Rome")); - assertThat(thrown).hasMessageThat().contains("tool node 'forecast'"); + assertThat(thrown).hasMessageThat().contains("confirmation request of node 'forecast'"); assertThat(thrown).hasMessageThat().contains("not a ToolConfirmation"); } @@ -649,6 +650,45 @@ public void resume_answerWithoutConfirmed_rejectsTheCall() { assertThat(database.deleted).isEmpty(); } + @Test + public void resume_answerWithConfirmedNull_rejectsTheCall() { + Database database = new Database(); + WorkflowSession session = new WorkflowSession(deletionWorkflow(database)); + session.advance(go()); + Map answer = new HashMap<>(); + answer.put("confirmed", null); + + ImmutableList events = session.turn(requestAnswer(DELETION_CONFIRMATION_ID, answer)); + + assertThat(replyInput(events)) + .isEqualTo(ImmutableMap.of("error", "This tool call is rejected.")); + assertThat(database.deleted).isEmpty(); + } + + @Test + public void resume_answerWithoutAHint_toolSeesAnEmptyHint() { + AtomicReference<@Nullable String> hint = new AtomicReference<>(); + FakeTool tool = + FakeTool.running( + "transfer", + (args, toolContext) -> { + if (toolContext.toolConfirmation().isEmpty()) { + toolContext.requestConfirmation("Approve the transfer?"); + return ImmutableMap.of(); + } + hint.set(toolContext.toolConfirmation().get().hint()); + return ImmutableMap.of("status", "sent"); + }); + WorkflowSession session = new WorkflowSession(toolWorkflow(tool, ImmutableMap.of("amount", 5))); + session.advance(go()); + + ImmutableList events = + session.turn(requestAnswer(TRANSFER_CONFIRMATION_ID, ImmutableMap.of("confirmed", true))); + + assertThat(replyInput(events)).isEqualTo(ImmutableMap.of("status", "sent")); + assertThat(hint.get()).isEmpty(); + } + @Test public void resume_confirmationApproved_showsTheAnswerToThePluginsBeforeToolCallback() { RecordingPlugin plugin = RecordingPlugin.silent();