fix: enforce node retry policies and timeouts

- Add `Context.attemptCount()`, which tells a workflow node which attempt it is on.
- When a failed child workflow runs again under the same run ID, rerun its nodes that failed because their own child failed.

PiperOrigin-RevId: 1008119395

Author: damianmomotgoogle

core/src/main/java/com/google/adk/agents/Context.java

@@ -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}.

core/src/main/java/com/google/adk/agents/InvocationContext.java

@@ -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();

core/src/main/java/com/google/adk/runner/Runner.java

@@ -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

core/src/main/java/com/google/adk/workflow/NodeConfig.java

@@ -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();
 

core/src/main/java/com/google/adk/workflow/NodeExecutionState.java

@@ -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.
@@ -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. */
@@ -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);
@@ -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);
+    }
   }
 
   /**
@@ -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}.
@@ -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() {}
 
@@ -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.
      *