Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
36 changes: 20 additions & 16 deletions core/src/main/java/com/google/adk/workflow/BaseNode.java
Original file line number Diff line number Diff line change
Expand Up @@ -132,22 +132,26 @@ protected void validateOutgoingEdges(Collection<Edge> edges) {}
*/
final Flowable<Event> 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.<Object>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("");
}
}
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
141 changes: 43 additions & 98 deletions core/src/main/java/com/google/adk/workflow/FunctionNode.java
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -56,25 +53,12 @@ public final class FunctionNode<I, O> extends BaseNode {
private static final PortCodec<Object> ANY =
PortCodec.output(PortCodec.javaType(Object.class), "an event output");

private static final String CONFIRMATION_INTERRUPT_PREFIX = "wf_tool_confirmation:";

private static final ConfirmationPredicate<Object> ALWAYS_CONFIRM = (context, input) -> true;

/** The output of a run the user rejected. */
private static final ImmutableMap<String, String> 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<I> inputCodec;
private final PortCodec<O> outputCodec;
private final NodeBody<? super I, ? extends O> body;
Expand Down Expand Up @@ -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<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,34 +185,51 @@ 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);
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<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()
: 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.
Expand All @@ -236,7 +240,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 Expand Up @@ -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<String, @Nullable Object> 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() + "'";
Expand Down
123 changes: 123 additions & 0 deletions core/src/main/java/com/google/adk/workflow/ToolConfirmations.java
Original file line number Diff line number Diff line change
@@ -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<String, @Nullable Object> 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() {}
}
Loading
Loading