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
89 changes: 80 additions & 9 deletions core/src/main/java/com/google/adk/workflow/FunctionNode.java
Original file line number Diff line number Diff line change
Expand Up @@ -23,10 +23,16 @@
import com.google.adk.agents.Context;
import com.google.adk.annotations.Experimental;
import com.google.adk.events.Event;
import com.google.common.base.Throwables;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
import com.google.genai.types.Schema;
import io.reactivex.rxjava3.core.Flowable;
import io.reactivex.rxjava3.core.Maybe;
import io.reactivex.rxjava3.core.MaybeEmitter;
import java.util.Optional;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.FutureTask;
import org.jspecify.annotations.Nullable;

/**
Expand Down Expand Up @@ -124,15 +130,42 @@ boolean edgeInputSchemaInferred() {

@Override
public Flowable<?> runNode(Context context, @Nullable Object nodeInput) {
return Flowable.defer(
() -> {
I input = inputCodec.decode(nodeInput, "input of node '" + name() + "'");
O result = body.apply(context, input);
encodeDelegatedOutput(context.requireNodeState());
// After a delegation, NodeRunner drops a returned output.
return Flowable.fromOptional(Optional.ofNullable(output(result)));
})
.subscribeOn(context.invocationContext().scheduler());
// Maybe.create rather than defer, so that canceling the node reaches the body while it runs.
return Maybe.<Object>create(emitter -> runBody(context, nodeInput, emitter))
.subscribeOn(context.invocationContext().scheduler())
.toFlowable();
}

/**
* Runs the body on the calling thread and sends its output or failure to {@code emitter}.
* Disposing {@code emitter} interrupts the body while it runs and drops its result.
*/
private void runBody(Context context, @Nullable Object nodeInput, MaybeEmitter<Object> emitter) {
Worker<O> worker =
new Worker<>(() -> {
I input = inputCodec.decode(nodeInput, "input of node '" + name() + "'");
return emitter.isDisposed() ? null : body.apply(context, input);
});
emitter.setCancellable(worker::cancel);
Object output;
try {
O result = worker.run();
if (emitter.isDisposed()) {
return;
}
encodeDelegatedOutput(context.requireNodeState());
// After a delegation, NodeRunner drops a returned output.
output = output(result);
} catch (Exception e) {
// tryOnError, as a canceled run can no longer take a failure.
emitter.tryOnError(e);
return;
}
if (output == null) {
emitter.onComplete();
} else {
emitter.onSuccess(output);
}
}

/**
Expand Down Expand Up @@ -188,6 +221,44 @@ private void encodeDelegatedOutput(NodeExecutionState state) {
return outputCodec.encode(result, what);
}

/**
* Runs one body on the thread that calls {@link #run}, and {@link #cancel} interrupts that thread
* only while the body runs. When the body returns, the worker clears its own interrupt, so the
* thread's next task does not see it. It keeps an interrupt the caller had set before the body
* ran. A cancel does not wait for a body that ignores the interrupt.
*/
private static final class Worker<T> {
private final FutureTask<T> task;

Worker(Callable<T> body) {
this.task = new FutureTask<>(body);
}

@SuppressWarnings("Interruption") // Interrupting the body is how a cancel reaches it.
void cancel() {
task.cancel(/* mayInterruptIfRunning= */ true);
}

/** Runs the body and returns its result, or {@code null} if it was canceled. */
@Nullable T run() throws Exception {
boolean callerInterrupted = Thread.currentThread().isInterrupted();
task.run();
if (task.isCancelled()) {
if (!callerInterrupted) {
var unused = Thread.interrupted();
}
return null;
}
try {
return task.get(); // Done, so get does not block.
} catch (ExecutionException e) {
Throwables.throwIfInstanceOf(e.getCause(), Exception.class);
Throwables.throwIfUnchecked(e.getCause());
throw e;
}
}
}

/**
* 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.
Expand Down
3 changes: 2 additions & 1 deletion core/src/main/java/com/google/adk/workflow/NodeBody.java
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,8 @@
* The body of a {@link FunctionNode}, which takes the node run's context and its single input; the
* input may be shared with other nodes, so the body must not mutate it. It runs on the {@code
* Runner}'s scheduler, so it may block, and a small fixed-size scheduler can deadlock bodies that
* wait for one another.
* wait for one another. Canceling the node interrupts the body's thread and drops its result; a
* body that ignores the interrupt keeps its thread until it returns, and a retry can overlap it.
*/
@Experimental
@FunctionalInterface
Expand Down
Loading
Loading