Skip to content
Merged
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
11 changes: 11 additions & 0 deletions core/src/main/java/com/google/adk/agents/Context.java
Original file line number Diff line number Diff line change
Expand Up @@ -334,6 +334,17 @@ public final void routes(List<? extends Route> routes) {
return requireNodeState().resumeInputs();
}

/**
* Returns the number of the attempt this workflow node is on, counting from 1; each retry adds
* one.
*
* @throws IllegalStateException if this is not the context of a workflow node.
*/
@Experimental
public final int attemptCount() {
return requireNodeState().attemptCount();
}

/**
* Runs {@code node} as a child of this workflow node without an input, as {@link #runNode(Node,
* Object, RunNodeConfig)} does with the default {@link RunNodeConfig}.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -314,9 +314,10 @@ public String userId() {
/**
* Returns the {@link Scheduler} on which this invocation runs the work ADK hands to worker
* threads: the sub-agents of a {@code ParallelAgent}, tool calls in {@code
* ToolExecutionMode.PARALLEL_SUBSCRIBE} mode, and the hop that starts the live send loop. An
* agent-level setting ({@code ParallelAgent.Builder#scheduler}, {@code
* LlmAgent.Builder#executor}) takes precedence over it. Defaults to {@link Schedulers#io()}.
* ToolExecutionMode.PARALLEL_SUBSCRIBE} mode, workflow node retries and attempts with a timeout,
* and the hop that starts the live send loop. An agent-level setting ({@code
* ParallelAgent.Builder#scheduler}, {@code LlmAgent.Builder#executor}) takes precedence over it.
* Defaults to {@link Schedulers#io()}.
*/
public Scheduler scheduler() {
return scheduler != null ? scheduler : Schedulers.io();
Expand Down
9 changes: 5 additions & 4 deletions core/src/main/java/com/google/adk/runner/Runner.java
Original file line number Diff line number Diff line change
Expand Up @@ -184,10 +184,11 @@ public Builder plugins(Plugin... plugins) {
/**
* Sets the {@link Scheduler} for the work ADK hands off to worker threads in this runner's
* invocations: the sub-agents of a {@code ParallelAgent}, tool calls in {@link
* RunConfig.ToolExecutionMode#PARALLEL_SUBSCRIBE} mode, and the hop that starts the live send
* loop. The runs that an {@code AgentTool} nests inherit it. An agent-level {@code
* ParallelAgent.Builder#scheduler} or {@code LlmAgent.Builder#executor} takes precedence.
* {@code null} (the default) selects {@code Schedulers.io()}.
* RunConfig.ToolExecutionMode#PARALLEL_SUBSCRIBE} mode, workflow node retries and attempts with
* a timeout, and the hop that starts the live send loop. The runs that an {@code AgentTool}
* nests inherit it. An agent-level {@code ParallelAgent.Builder#scheduler} or {@code
* LlmAgent.Builder#executor} takes precedence. {@code null} (the default) selects {@code
* Schedulers.io()}.
*
* <p>Only these hand-offs are affected; work that completes asynchronously elsewhere (model
* clients, session services, MCP, plugins) stays on its own threads. {@code
Expand Down
13 changes: 10 additions & 3 deletions core/src/main/java/com/google/adk/workflow/NodeConfig.java
Original file line number Diff line number Diff line change
Expand Up @@ -18,24 +18,31 @@

import static com.google.common.base.Preconditions.checkArgument;

import com.google.adk.agents.Context;
import com.google.adk.annotations.Experimental;
import com.google.auto.value.AutoValue;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
import java.time.Duration;
import java.util.Optional;
import java.util.concurrent.TimeoutException;
import org.jspecify.annotations.Nullable;

/** A node's retry policy and execution timeout. */
@Experimental
@AutoValue
public abstract class NodeConfig {

/** Returns how the node is retried when it fails; empty does not retry. */
/**
* Returns how the node is retried when it fails; empty does not retry. A retried workflow replays
* its nodes that finished and left an event in the session, and reruns the others. A failure that
* a node passes on unchanged from a child it runs with {@link Context#runNode} is not retried;
* give the child a retry policy of its own instead.
*/
public abstract Optional<RetryConfig> retryConfig();

/**
* Returns how long an attempt may run before it is cancelled and treated as a failure; empty
* imposes no limit.
* Returns how long an attempt may run before it is canceled and fails with a {@link
* TimeoutException}; empty imposes no limit.
*/
public abstract Optional<Duration> timeout();

Expand Down
51 changes: 47 additions & 4 deletions core/src/main/java/com/google/adk/workflow/NodeExecutionState.java
Original file line number Diff line number Diff line change
Expand Up @@ -61,13 +61,16 @@ public final class NodeExecutionState {
// The parent's state when this run's output is also the parent's.
private final @Nullable NodeExecutionState outputParent;
private final String runId;
private final int attemptCount;
private final String nodePath;
private final ImmutableList<String> outputFor;
private final Map<String, @Nullable Object> resumeInputs;
private final ConcurrentMap<String, Integer> childRunCounts = new ConcurrentHashMap<>();
private final ConcurrentMap<String, SingleSubject<Context>> childRuns = new ConcurrentHashMap<>();
// The paths of the node runs started but not finished in this turn, shared under one root.
private final ConcurrentMap<String, Integer> childRunCounts;
private final ConcurrentMap<String, SingleSubject<Context>> childRuns;
// Paths of runs started but not finished in this turn, shared under one root and its attempts.
private final Set<String> unfinishedNodePaths;
// Paths of runs that failed this turn and have not finished cleanly since, shared the same way.
private final Set<String> failedNodePaths;
private final Set<String> interruptIds = new CopyOnWriteArraySet<>();
private final AtomicReference<OutputState> outputState = new AtomicReference<>(NO_OUTPUT);
// Set once the node drops its own output because a useAsOutput child holds the output.
Expand Down Expand Up @@ -100,9 +103,19 @@ private NodeExecutionState(Builder builder) {
? ImmutableList.<String>builder().add(nodePath).addAll(outputParent.outputFor).build()
: ImmutableList.of(nodePath);
this.eventAuthor = parentState == null ? "" : parentState.eventAuthor;
NodeExecutionState previousAttempt = builder.previousAttempt;
this.attemptCount = previousAttempt == null ? 1 : previousAttempt.attemptCount + 1;
NodeExecutionState turnRecords = parentState != null ? parentState : previousAttempt;
this.unfinishedNodePaths =
parentState == null ? ConcurrentHashMap.newKeySet() : parentState.unfinishedNodePaths;
turnRecords == null ? ConcurrentHashMap.newKeySet() : turnRecords.unfinishedNodePaths;
this.failedNodePaths =
turnRecords == null ? ConcurrentHashMap.newKeySet() : turnRecords.failedNodePaths;
this.resumeInputs = Collections.unmodifiableMap(new LinkedHashMap<>(builder.resumeInputs));
// A retried workflow restarts its node numbering, so its finished nodes replay.
boolean continuesChildRuns = previousAttempt != null && !(node instanceof Workflow);
this.childRunCounts =
continuesChildRuns ? previousAttempt.childRunCounts : new ConcurrentHashMap<>();
this.childRuns = continuesChildRuns ? previousAttempt.childRuns : new ConcurrentHashMap<>();
}

/** Returns a builder for the state of one node run. */
Expand All @@ -128,6 +141,11 @@ public String runId() {
return runId;
}

/** Returns the number of this attempt at running the node, counting from 1. */
public int attemptCount() {
return attemptCount;
}

/** Returns the context of the node that ran this node, or empty if this is a root node. */
public Optional<Context> parent() {
return Optional.ofNullable(parent);
Expand Down Expand Up @@ -266,6 +284,11 @@ void markStarted() {
/** Records that this run finished, with or without a result. */
void markFinished() {
unfinishedNodePaths.remove(nodePath);
if (failure == null) {
failedNodePaths.remove(nodePath);
} else {
failedNodePaths.add(nodePath);
}
}

/**
Expand All @@ -277,6 +300,14 @@ boolean childCanceledMidRun(String childKey) {
return unfinishedNodePaths.contains(nodePath + "/" + childKey);
}

/**
* Returns whether this run's child {@code childKey} failed in this turn and has not finished
* cleanly since. A child that takes on its own child's failure leaves no error event to show it.
*/
boolean childFailedThisTurn(String childKey) {
return failedNodePaths.contains(nodePath + "/" + childKey);
}

/**
* Returns the states of this node's children from earlier turns, keyed by {@code name@runId},
* which the first call reconstructs from the events of {@code invocationContext}.
Expand Down Expand Up @@ -457,6 +488,7 @@ static final class Builder {
private @Nullable String runId;
private boolean useAsOutput;
private Map<String, @Nullable Object> resumeInputs = ImmutableMap.of();
private @Nullable NodeExecutionState previousAttempt;

private Builder() {}

Expand Down Expand Up @@ -505,6 +537,17 @@ Builder resumeInputs(Map<String, @Nullable Object> resumeInputs) {
return this;
}

/**
* Sets the state of the run's previous attempt, or null for a first attempt. A later attempt's
* count is one higher. It shares the previous attempt's records of unfinished and failed runs
* and, unless the node is a workflow, its child run IDs and child runs.
*/
@CanIgnoreReturnValue
Builder previousAttempt(@Nullable NodeExecutionState previousAttempt) {
this.previousAttempt = previousAttempt;
return this;
}

/**
* Builds the state.
*
Expand Down
Loading
Loading