From 367b402db39f93ab430f473f48bce89a301bbb60 Mon Sep 17 00:00:00 2001 From: Damian Momot Date: Tue, 27 Oct 2026 08:21:28 -0700 Subject: [PATCH] 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 --- .../java/com/google/adk/agents/Context.java | 11 + .../google/adk/agents/InvocationContext.java | 7 +- .../java/com/google/adk/runner/Runner.java | 9 +- .../com/google/adk/workflow/NodeConfig.java | 13 +- .../adk/workflow/NodeExecutionState.java | 51 +- .../com/google/adk/workflow/NodeRunner.java | 142 +++- .../com/google/adk/workflow/ResumeScan.java | 18 +- .../com/google/adk/workflow/RetryConfig.java | 39 + .../com/google/adk/workflow/Scheduler.java | 3 +- .../com/google/adk/agents/ContextTest.java | 1 + .../google/adk/workflow/ResumeScanTest.java | 80 +- .../google/adk/workflow/RetryConfigTest.java | 114 +++ .../adk/workflow/WorkflowResumeTest.java | 16 +- .../adk/workflow/WorkflowResumeTesting.java | 14 + .../adk/workflow/WorkflowRetryTest.java | 771 ++++++++++++++++++ 15 files changed, 1215 insertions(+), 74 deletions(-) create mode 100644 core/src/test/java/com/google/adk/workflow/WorkflowRetryTest.java diff --git a/core/src/main/java/com/google/adk/agents/Context.java b/core/src/main/java/com/google/adk/agents/Context.java index b6d1991c3..a046498dd 100644 --- a/core/src/main/java/com/google/adk/agents/Context.java +++ b/core/src/main/java/com/google/adk/agents/Context.java @@ -334,6 +334,17 @@ public final void routes(List 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}. diff --git a/core/src/main/java/com/google/adk/agents/InvocationContext.java b/core/src/main/java/com/google/adk/agents/InvocationContext.java index 3a761dedc..338eaa00f 100644 --- a/core/src/main/java/com/google/adk/agents/InvocationContext.java +++ b/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(); diff --git a/core/src/main/java/com/google/adk/runner/Runner.java b/core/src/main/java/com/google/adk/runner/Runner.java index d80f60d4b..8c62d14ce 100644 --- a/core/src/main/java/com/google/adk/runner/Runner.java +++ b/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()}. * *

Only these hand-offs are affected; work that completes asynchronously elsewhere (model * clients, session services, MCP, plugins) stays on its own threads. {@code diff --git a/core/src/main/java/com/google/adk/workflow/NodeConfig.java b/core/src/main/java/com/google/adk/workflow/NodeConfig.java index b397442f0..e17420e41 100644 --- a/core/src/main/java/com/google/adk/workflow/NodeConfig.java +++ b/core/src/main/java/com/google/adk/workflow/NodeConfig.java @@ -18,11 +18,13 @@ 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. */ @@ -30,12 +32,17 @@ @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(); /** - * 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 timeout(); diff --git a/core/src/main/java/com/google/adk/workflow/NodeExecutionState.java b/core/src/main/java/com/google/adk/workflow/NodeExecutionState.java index 4bf0762ca..b4a557b1b 100644 --- a/core/src/main/java/com/google/adk/workflow/NodeExecutionState.java +++ b/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 outputFor; private final Map resumeInputs; - private final ConcurrentMap childRunCounts = new ConcurrentHashMap<>(); - private final ConcurrentMap> childRuns = new ConcurrentHashMap<>(); - // The paths of the node runs started but not finished in this turn, shared under one root. + private final ConcurrentMap childRunCounts; + private final ConcurrentMap> childRuns; + // Paths of runs started but not finished in this turn, shared under one root and its attempts. private final Set unfinishedNodePaths; + // Paths of runs that failed this turn and have not finished cleanly since, shared the same way. + private final Set failedNodePaths; private final Set interruptIds = new CopyOnWriteArraySet<>(); private final AtomicReference 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.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 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 resumeInputs = ImmutableMap.of(); + private @Nullable NodeExecutionState previousAttempt; private Builder() {} @@ -505,6 +537,17 @@ Builder resumeInputs(Map 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. * diff --git a/core/src/main/java/com/google/adk/workflow/NodeRunner.java b/core/src/main/java/com/google/adk/workflow/NodeRunner.java index 83b27351f..e7b0258f8 100644 --- a/core/src/main/java/com/google/adk/workflow/NodeRunner.java +++ b/core/src/main/java/com/google/adk/workflow/NodeRunner.java @@ -21,12 +21,14 @@ import static com.google.common.base.Preconditions.checkState; import static com.google.common.base.Strings.isNullOrEmpty; import static com.google.common.base.Strings.nullToEmpty; +import static java.util.concurrent.TimeUnit.NANOSECONDS; 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.telemetry.Tracing; import com.google.adk.workflow.ResumeScan.Interception; import com.google.common.base.CharMatcher; import com.google.common.collect.ImmutableMap; @@ -41,19 +43,25 @@ import io.reactivex.rxjava3.processors.UnicastProcessor; import io.reactivex.rxjava3.subjects.CompletableSubject; import io.reactivex.rxjava3.subjects.SingleSubject; +import java.time.Duration; import java.util.Map; import java.util.Optional; import java.util.Queue; import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.ThreadLocalRandom; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicReference; import java.util.function.Function; import org.jspecify.annotations.Nullable; /** - * The runner for one workflow node. It stamps every emitted event with its author, node path, and - * branch before forwarding it, and records the node's output, routes, interrupts, and failure on - * the node state of the returned context. {@link #runRoot} runs a node as the root of an - * invocation, and {@link #runChild} runs one as a child that another node starts itself. + * The runner for one workflow node, which times out an attempt that exceeds the node's optional + * timeout and retries the node as its retry policy allows. It stamps every emitted event with its + * author, node path, and branch before forwarding it, and records the node's output, routes, + * interrupts, and failure on the node state of the returned context. {@link #runRoot} runs a node + * as the root of an invocation, and {@link #runChild} runs one as a child that another node starts + * itself. */ @Experimental public final class NodeRunner { @@ -286,7 +294,11 @@ private static Single runOrReplayChild( parentState.recoveredChildren(parent.invocationContext()).get(childKey); Interception interception = ResumeScan.intercept( - node, recovered, /* dynamic= */ true, parentState.childCanceledMidRun(childKey)); + node, + recovered, + /* dynamic= */ true, + parentState.childCanceledMidRun(childKey), + parentState.childFailedThisTurn(childKey)); if (!interception.shouldRun()) { return Single.just( ResumeScan.replayContext(parent, node, runId, config.useAsOutput(), interception)); @@ -364,16 +376,15 @@ private static Maybe childResult( } /** - * Runs the node on {@code nodeInput} and returns its context when the node finishes. An exception - * from the node is recorded on the context's node state rather than signaled, while an {@code - * Error} propagates. + * Runs the node on {@code nodeInput}, retrying it as its retry policy allows, and returns the + * context of its last attempt. An exception from the node is recorded on the context's node state + * rather than signaled, while an {@code Error} propagates. */ Single run(@Nullable Object nodeInput) { return Single.defer( () -> { if (parent != null) { - Context context = newContext(resumeInputs); - return new Activation(context).run(nodeInput).toSingleDefault(context); + return runAttempts(nodeInput, resumeInputs); } // A root scans its own progress and resumes like a dynamic child. Interception interception = @@ -382,17 +393,55 @@ Single run(@Nullable Object nodeInput) { ResumeScan.scan(invocationContext, /* basePath= */ "") .get(segment(node.name(), runId)), /* dynamic= */ true, - /* canceledMidRun= */ false); - Context context = newContext(interception.resumeInputs()); + /* canceledMidRun= */ false, + /* failedThisTurn= */ false); if (!interception.shouldRun()) { + Context context = newContext(interception.resumeInputs(), /* previousAttempt= */ null); interception.replayOn(context); return new Activation(context).flushPending().toSingleDefault(context); } - return new Activation(context).run(nodeInput).toSingleDefault(context); + return runAttempts(nodeInput, interception.resumeInputs()); }); } - private Context newContext(Map resumeInputs) { + private Single runAttempts( + @Nullable Object nodeInput, Map resumeInputs) { + RetryConfig policy = node.config().retryConfig().orElse(null); + int maxAttempts = policy == null ? 1 : Math.max(1, policy.maxAttempts()); + AtomicReference<@Nullable NodeExecutionState> previousAttempt = new AtomicReference<>(); + // A retry starts on another thread, so each attempt restores the caller's trace context. + io.opentelemetry.context.Context traceContext = io.opentelemetry.context.Context.current(); + // concatMap loops instead of recursing, so many instant retries keep the stack flat. + return Flowable.range(1, maxAttempts) + .concatMapSingle( + attempt -> + backoff(policy, attempt) + .andThen( + Single.defer( + () -> { + Context context = newContext(resumeInputs, previousAttempt.get()); + previousAttempt.set(context.requireNodeState()); + Activation activation = new Activation(context); + return activation.run(nodeInput).toSingleDefault(activation); + }) + .compose(Tracing.withContext(traceContext)))) + .takeUntil(activation -> !activation.retries(policy)) + .lastOrError() + .map(activation -> activation.context); + } + + /** Returns the wait before attempt {@code attempt}; there is none before the first. */ + private Completable backoff(@Nullable RetryConfig policy, int attempt) { + if (policy == null || attempt == 1) { + return Completable.complete(); + } + Duration delay = policy.delayFor(attempt - 1, ThreadLocalRandom.current()); + // Even without a delay, the retry starts on the invocation's scheduler, not the failing thread. + return Completable.timer(delay.toNanos(), NANOSECONDS, invocationContext.scheduler()); + } + + private Context newContext( + Map resumeInputs, @Nullable NodeExecutionState previousAttempt) { return new Context( childInvocationContext(), NodeExecutionState.builder() @@ -402,6 +451,7 @@ private Context newContext(Map resumeInputs) { .runId(runId) .useAsOutput(useAsOutput) .resumeInputs(resumeInputs) + .previousAttempt(previousAttempt) .build()); } @@ -420,14 +470,16 @@ private InvocationContext childInvocationContext() { } /** - * One run of the node, during which the node is asked for each event only after its previous - * event is processed. Events from nested nodes pass through without reapplying their routing or - * transfers. When the node finishes, a final event carries any output, routes, or deltas not - * carried by an earlier event. + * One attempt at running the node, during which the node is asked for each event only after its + * previous event is processed. Events from nested nodes pass through without reapplying their + * routing or transfers. When the node finishes, a final event carries any output, routes, or + * deltas not carried by an earlier event. */ private final class Activation { private final Context context; private final NodeExecutionState nodeState; + // Set when the node takes on a child's failure, which only the child's own policy retries. + private boolean childFailed; Activation(Context context) { this.context = context; @@ -435,10 +487,12 @@ private final class Activation { } Completable run(@Nullable Object nodeInput) { - return asBaseNode(node) - .run(context, nodeInput) - // Avoid concatMapCompletable: its prefetch would let the node run ahead of the sink. - .flatMapCompletable(this::dispatch, /* delayErrors= */ true, /* maxConcurrency= */ 1) + Completable body = + asBaseNode(node) + .run(context, nodeInput) + // Avoid concatMapCompletable: its prefetch would let the node run ahead of the sink. + .flatMapCompletable(this::dispatch, /* delayErrors= */ true, /* maxConcurrency= */ 1); + return withTimeout(body) .andThen(Completable.defer(this::flushPending)) .onErrorResumeNext(this::onError) // A run cut short stays unfinished, so a later call in this turn reruns it. @@ -446,6 +500,41 @@ Completable run(@Nullable Object nodeInput) { .doOnComplete(nodeState::markFinished); } + /** Returns whether {@code policy} retries this attempt. */ + boolean retries(@Nullable RetryConfig policy) { + return policy != null + && !childFailed + && nodeState + .failure() + .map(failure -> policy.shouldRetry(failure, nodeState.attemptCount())) + .orElse(false); + } + + /** Fails {@code run} with a {@link TimeoutException} once it outlasts the node's timeout. */ + private Completable withTimeout(Completable run) { + Optional timeout = node.config().timeout(); + if (timeout.isEmpty()) { + return run; + } + Duration limit = timeout.get(); + // Fails on the timer's thread: the timed-out attempt may still hold the scheduler's threads. + Completable timedOut = + Completable.error( + () -> + new TimeoutException( + String.format( + "Node '%s' timed out after %d ms.", node.name(), limit.toMillis()))); + io.opentelemetry.context.Context traceContext = io.opentelemetry.context.Context.current(); + // Off the caller's thread, so a body that blocks its thread cannot hold off the timeout. + return Completable.timer(0, NANOSECONDS, invocationContext.scheduler()) + // Unlike subscribeOn, canceling this hop does not interrupt the handling of node events. + .andThen(run.compose(Tracing.withContext(traceContext))) + // convert saturates, so a timeout too long to count in nanoseconds waits about 292 years. + .timeout(NANOSECONDS.convert(limit), NANOSECONDS, timedOut) + // The failure arrives on the timer's thread, so it needs the caller's trace context too. + .compose(Tracing.withContext(traceContext)); + } + private Completable onError(Throwable error) { if (error instanceof NodeInterruptedException) { // A child waits, or raiseOnWait ended the run; any interrupt IDs are already recorded. @@ -453,6 +542,7 @@ private Completable onError(Throwable error) { } if (error instanceof DynamicNodeFailedException childFailure) { // The child already sent an error event, so the node takes on its failure without one. + childFailed = true; nodeState.failure(childFailure.error(), childFailure.errorNodePath()); return Completable.complete(); } @@ -576,7 +666,7 @@ private Completable fail(Exception failure) { withPendingDeltas( stamp( newEvent() - .errorCode(new FinishReason(errorTypeName(failure))) + .errorCode(new FinishReason(RetryConfig.errorTypeName(failure))) .errorMessage(nullToEmpty(failure.getMessage())) .build())); return eventSink.send(errorEvent, failure).doOnComplete(() -> nodeState.failure(failure)); @@ -657,12 +747,6 @@ private static boolean isMessageAsOutput(Event event) { && event.nodeInfo().map(NodeInfo::messageAsOutput).orElse(false); } - /** Returns the simple name of the failure's type, or its full name for an anonymous class. */ - private static String errorTypeName(Exception failure) { - String name = failure.getClass().getSimpleName(); - return name.isEmpty() ? failure.getClass().getName() : name; - } - /** Adapts any node to a {@link BaseNode}, whose {@link BaseNode#run} turns values into events. */ private static BaseNode asBaseNode(Node node) { if (node instanceof BaseNode baseNode) { diff --git a/core/src/main/java/com/google/adk/workflow/ResumeScan.java b/core/src/main/java/com/google/adk/workflow/ResumeScan.java index 7d8ca72c5..28dbf9dfc 100644 --- a/core/src/main/java/com/google/adk/workflow/ResumeScan.java +++ b/core/src/main/java/com/google/adk/workflow/ResumeScan.java @@ -105,13 +105,18 @@ static ImmutableMap scan( /** * Decides whether {@code node} runs again or replays its result from {@code recovered}. In order: * a node with unanswered interrupts waits on them, unless it reruns on resume and has answers; a - * failed node reruns; a recorded output, route or transfer replays; a node {@code canceledMidRun} - * in this turn reruns; an answered node that does not rerun on resume completes with its answers. - * Otherwise a node that reruns on resume reruns unless it finished since its last answer, and a - * {@code dynamic} node reruns while a graph node does not. + * node that failed, in {@code recovered} or {@code failedThisTurn}, reruns; a recorded output, + * route or transfer replays; a node {@code canceledMidRun} in this turn reruns; an answered node + * that does not rerun on resume completes with its answers. Otherwise a node that reruns on + * resume reruns unless it finished since its last answer, and a {@code dynamic} node reruns while + * a graph node does not. */ static Interception intercept( - Node node, @Nullable RecoveredNode recovered, boolean dynamic, boolean canceledMidRun) { + Node node, + @Nullable RecoveredNode recovered, + boolean dynamic, + boolean canceledMidRun, + boolean failedThisTurn) { if (recovered == null) { return Interception.run(ImmutableMap.of()); } @@ -129,7 +134,8 @@ static Interception intercept( ? Interception.run(resolved) : Interception.waitOn(branch, unresolved, transfer); } - if (recovered.failed) { + // A node that takes on a child's failure leaves no error event, so this turn records it. + if (recovered.failed || failedThisTurn) { return Interception.run(resolved); } Object output = recovered.output; diff --git a/core/src/main/java/com/google/adk/workflow/RetryConfig.java b/core/src/main/java/com/google/adk/workflow/RetryConfig.java index ce79fa622..451d2f3eb 100644 --- a/core/src/main/java/com/google/adk/workflow/RetryConfig.java +++ b/core/src/main/java/com/google/adk/workflow/RetryConfig.java @@ -17,6 +17,7 @@ package com.google.adk.workflow; import static com.google.common.base.Preconditions.checkArgument; +import static java.util.concurrent.TimeUnit.NANOSECONDS; import com.google.adk.annotations.Experimental; import com.google.auto.value.AutoValue; @@ -25,6 +26,7 @@ import java.time.Duration; import java.util.List; import java.util.Optional; +import java.util.Random; import org.jspecify.annotations.Nullable; /** @@ -85,6 +87,43 @@ public static Builder builder() { public abstract Builder toBuilder(); + /** + * Returns whether a node that failed with {@code failure} on attempt {@code attempt}, counting + * from 1, has an attempt left and failed with a type this config retries. + */ + boolean shouldRetry(Exception failure, int attempt) { + if (attempt >= maxAttempts()) { + return false; + } + return exceptions().map(names -> names.contains(errorTypeName(failure))).orElse(true); + } + + /** + * Returns how long to wait after the failed attempt {@code attempt}, counting from 1, before the + * next one, drawing the jitter from {@code random}. + */ + Duration delayFor(int attempt, Random random) { + // TimeUnit.convert saturates instead of overflowing past about 292 years. + double maxNanos = NANOSECONDS.convert(maxDelay()); + // Clamped, as an infinite power times a zero initial delay would be NaN. + double delay = + NANOSECONDS.convert(initialDelay()) + * Math.min(Math.pow(backoffFactor(), attempt - 1), Double.MAX_VALUE); + // Capped before jittering, so jitter keeps spreading the delays near the cap. + double capped = Math.min(delay, maxNanos / (1.0 + jitter())); + double spread = jitter() * (2.0 * random.nextDouble() - 1.0); + return Duration.ofNanos((long) Math.max(0.0, Math.min(capped * (1.0 + spread), maxNanos))); + } + + /** + * Returns the name that a node's error event reports and {@link #exceptions} matches: the simple + * name of the failure's type, or its full name for an anonymous class. + */ + static String errorTypeName(Exception failure) { + String name = failure.getClass().getSimpleName(); + return name.isEmpty() ? failure.getClass().getName() : name; + } + /** Builder for {@link RetryConfig}. */ @AutoValue.Builder public abstract static class Builder { diff --git a/core/src/main/java/com/google/adk/workflow/Scheduler.java b/core/src/main/java/com/google/adk/workflow/Scheduler.java index 08e7737f1..09a116948 100644 --- a/core/src/main/java/com/google/adk/workflow/Scheduler.java +++ b/core/src/main/java/com/google/adk/workflow/Scheduler.java @@ -211,7 +211,8 @@ private void start(String nodeName, Trigger trigger) { node, recoveredNode, /* dynamic= */ false, - workflowState.childCanceledMidRun(childKey)); + workflowState.childCanceledMidRun(childKey), + workflowState.childFailedThisTurn(childKey)); if (recoveredNode != null && !interception.shouldRun()) { // Only a replay that completes an answered transfer marks a new step, as a run does. boolean completesTransfer = diff --git a/core/src/test/java/com/google/adk/agents/ContextTest.java b/core/src/test/java/com/google/adk/agents/ContextTest.java index e11549055..1075a0f00 100644 --- a/core/src/test/java/com/google/adk/agents/ContextTest.java +++ b/core/src/test/java/com/google/adk/agents/ContextTest.java @@ -95,6 +95,7 @@ public void nodeMethods_onCallbackContext_throwIllegalStateException() { assertThrows(IllegalStateException.class, context::output); assertThrows(IllegalStateException.class, context::routes); assertThrows(IllegalStateException.class, () -> context.routes(ImmutableList.of())); + assertThrows(IllegalStateException.class, context::attemptCount); } @Test diff --git a/core/src/test/java/com/google/adk/workflow/ResumeScanTest.java b/core/src/test/java/com/google/adk/workflow/ResumeScanTest.java index 690b681d7..36308a557 100644 --- a/core/src/test/java/com/google/adk/workflow/ResumeScanTest.java +++ b/core/src/test/java/com/google/adk/workflow/ResumeScanTest.java @@ -132,7 +132,11 @@ public void scan_outputOnAnEventWithAnErrorCode_replaysTheOutput() { RecoveredNode recovered = ResumeScan.scan(events, "wf@1", "inv").get("chat@1"); Interception interception = ResumeScan.intercept( - emitter("chat", "fresh"), recovered, /* dynamic= */ false, /* canceledMidRun= */ false); + emitter("chat", "fresh"), + recovered, + /* dynamic= */ false, + /* canceledMidRun= */ false, + /* failedThisTurn= */ false); assertThat(recovered.failed).isFalse(); assertThat(interception.shouldRun()).isFalse(); @@ -211,9 +215,19 @@ public void intercept_finishedRerunNodeThatWaitsForOutput_stillRuns() { unasked.finishedAfterResume = true; Interception answeredInterception = - ResumeScan.intercept(node, answered, /* dynamic= */ false, /* canceledMidRun= */ false); + ResumeScan.intercept( + node, + answered, + /* dynamic= */ false, + /* canceledMidRun= */ false, + /* failedThisTurn= */ false); Interception unaskedInterception = - ResumeScan.intercept(node, unasked, /* dynamic= */ false, /* canceledMidRun= */ false); + ResumeScan.intercept( + node, + unasked, + /* dynamic= */ false, + /* canceledMidRun= */ false, + /* failedThisTurn= */ false); assertThat(answeredInterception.shouldRun()).isTrue(); assertThat(unaskedInterception.shouldRun()).isTrue(); @@ -225,9 +239,19 @@ public void intercept_nothingRecorded_runsADynamicNodeButSkipsAGraphNode() { RecoveredNode recovered = new RecoveredNode(); Interception dynamic = - ResumeScan.intercept(node, recovered, /* dynamic= */ true, /* canceledMidRun= */ false); + ResumeScan.intercept( + node, + recovered, + /* dynamic= */ true, + /* canceledMidRun= */ false, + /* failedThisTurn= */ false); Interception graph = - ResumeScan.intercept(node, recovered, /* dynamic= */ false, /* canceledMidRun= */ false); + ResumeScan.intercept( + node, + recovered, + /* dynamic= */ false, + /* canceledMidRun= */ false, + /* failedThisTurn= */ false); assertThat(dynamic.shouldRun()).isTrue(); assertThat(graph.shouldRun()).isFalse(); @@ -239,7 +263,27 @@ public void intercept_nodeCanceledMidRun_runsAgainAsAGraphNode() { Interception interception = ResumeScan.intercept( - emitter("n", "out"), progressOnly, /* dynamic= */ false, /* canceledMidRun= */ true); + emitter("n", "out"), + progressOnly, + /* dynamic= */ false, + /* canceledMidRun= */ true, + /* failedThisTurn= */ false); + + assertThat(interception.shouldRun()).isTrue(); + } + + @Test + public void intercept_nodeThatFailedThisTurn_runsAgainDespiteARecordedOutput() { + RecoveredNode recovered = new RecoveredNode(); + recovered.output = "out"; + + Interception interception = + ResumeScan.intercept( + emitter("n", "out"), + recovered, + /* dynamic= */ false, + /* canceledMidRun= */ false, + /* failedThisTurn= */ true); assertThat(interception.shouldRun()).isTrue(); } @@ -253,7 +297,11 @@ public void intercept_answeredNodeThatDoesNotRerun_completesWithAllAnswersAsAMap Interception interception = ResumeScan.intercept( - emitter("n", "out"), recovered, /* dynamic= */ false, /* canceledMidRun= */ false); + emitter("n", "out"), + recovered, + /* dynamic= */ false, + /* canceledMidRun= */ false, + /* failedThisTurn= */ false); assertThat(interception.shouldRun()).isFalse(); assertThat(interception.output()).isEqualTo(ImmutableMap.of("a", 1, "b", 2)); @@ -271,13 +319,25 @@ public void intercept_unresolvedInterruptWithPartialProgress_runsOnlyANodeThatRe Interception rerunWithProgress = ResumeScan.intercept( - rerunNode, partialProgress, /* dynamic= */ false, /* canceledMidRun= */ false); + rerunNode, + partialProgress, + /* dynamic= */ false, + /* canceledMidRun= */ false, + /* failedThisTurn= */ false); Interception rerunWithoutProgress = ResumeScan.intercept( - rerunNode, noProgress, /* dynamic= */ false, /* canceledMidRun= */ false); + rerunNode, + noProgress, + /* dynamic= */ false, + /* canceledMidRun= */ false, + /* failedThisTurn= */ false); Interception staticWithProgress = ResumeScan.intercept( - graphNode, partialProgress, /* dynamic= */ false, /* canceledMidRun= */ false); + graphNode, + partialProgress, + /* dynamic= */ false, + /* canceledMidRun= */ false, + /* failedThisTurn= */ false); assertThat(rerunWithProgress.shouldRun()).isTrue(); assertThat(rerunWithProgress.resumeInputs()).containsExactly("a", 1); diff --git a/core/src/test/java/com/google/adk/workflow/RetryConfigTest.java b/core/src/test/java/com/google/adk/workflow/RetryConfigTest.java index 99b03b298..c563ff6db 100644 --- a/core/src/test/java/com/google/adk/workflow/RetryConfigTest.java +++ b/core/src/test/java/com/google/adk/workflow/RetryConfigTest.java @@ -16,13 +16,18 @@ package com.google.adk.workflow; +import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertThrows; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; import java.time.Duration; import java.util.ArrayList; +import java.util.Collections; import java.util.List; +import java.util.Random; +import java.util.stream.Stream; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -142,4 +147,113 @@ public void build_rejectsNaNJitter() { assertThat(e).hasMessageThat().contains("jitter"); } + + @Test + public void shouldRetry_withAttemptsLeftAndExceptionsUnset_retriesAnyFailure() { + RetryConfig config = RetryConfig.builder().maxAttempts(3).build(); + + assertThat(config.shouldRetry(new IllegalStateException(), /* attempt= */ 1)).isTrue(); + assertThat(config.shouldRetry(new IllegalArgumentException(), /* attempt= */ 2)).isTrue(); + } + + @Test + public void shouldRetry_withNoAttemptLeft_stops() { + Exception failure = new IllegalStateException(); + + assertThat(withMaxAttempts(3).shouldRetry(failure, /* attempt= */ 3)).isFalse(); + assertThat(withMaxAttempts(1).shouldRetry(failure, /* attempt= */ 1)).isFalse(); + assertThat(withMaxAttempts(0).shouldRetry(failure, /* attempt= */ 1)).isFalse(); + } + + @Test + public void shouldRetry_matchesTheSimpleTypeNameWithoutSubclasses() { + RetryConfig config = + RetryConfig.builder().exceptions(ImmutableList.of("RuntimeException")).build(); + + assertThat(config.shouldRetry(new RuntimeException(), /* attempt= */ 1)).isTrue(); + assertThat(config.shouldRetry(new IllegalStateException(), /* attempt= */ 1)).isFalse(); + } + + @Test + public void shouldRetry_withAnEmptyExceptionsList_retriesNothing() { + RetryConfig config = RetryConfig.builder().exceptions(ImmutableList.of()).build(); + + assertThat(config.shouldRetry(new IllegalStateException(), /* attempt= */ 1)).isFalse(); + } + + @Test + public void delayFor_withoutJitter_growsByTheBackoffFactorUntilMaxDelay() { + RetryConfig config = + RetryConfig.builder() + .initialDelay(Duration.ofSeconds(1)) + .maxDelay(Duration.ofSeconds(60)) + .backoffFactor(2) + .jitter(0) + .build(); + Random random = new Random(0); + + assertThat(config.delayFor(/* attempt= */ 1, random)).isEqualTo(Duration.ofSeconds(1)); + assertThat(config.delayFor(/* attempt= */ 2, random)).isEqualTo(Duration.ofSeconds(2)); + assertThat(config.delayFor(/* attempt= */ 3, random)).isEqualTo(Duration.ofSeconds(4)); + assertThat(config.delayFor(/* attempt= */ 10, random)).isEqualTo(Duration.ofSeconds(60)); + } + + @Test + public void delayFor_withJitter_spreadsTheDelayAroundItsBase() { + RetryConfig config = + RetryConfig.builder().initialDelay(Duration.ofSeconds(1)).jitter(0.5).build(); + Random random = new Random(0); + + ImmutableList delays = + Stream.generate(() -> config.delayFor(/* attempt= */ 1, random)) + .limit(50) + .collect(toImmutableList()); + + assertThat(Collections.min(delays)).isAtLeast(Duration.ofMillis(500)); + assertThat(Collections.max(delays)).isAtMost(Duration.ofMillis(1500)); + assertThat(ImmutableSet.copyOf(delays).size()).isGreaterThan(1); + } + + @Test + public void delayFor_withJitterWhenTheBackoffOverflows_staysWithinMaxDelay() { + RetryConfig config = + RetryConfig.builder() + .initialDelay(Duration.ofSeconds(1)) + .maxDelay(Duration.ofSeconds(3)) + .backoffFactor(2) + .jitter(0.5) + .build(); + Random random = new Random(0); + + ImmutableList delays = + Stream.generate(() -> config.delayFor(/* attempt= */ 5000, random)) + .limit(50) + .collect(toImmutableList()); + + // The delay is capped at maxDelay / (1 + jitter) = 2s before jitter spreads it. + assertThat(Collections.min(delays)).isAtLeast(Duration.ofSeconds(1)); + assertThat(Collections.max(delays)).isAtMost(Duration.ofSeconds(3)); + } + + @Test + public void delayFor_withAZeroInitialDelay_staysZeroWhenTheBackoffOverflows() { + RetryConfig config = + RetryConfig.builder().initialDelay(Duration.ZERO).backoffFactor(10).build(); + + assertThat(config.delayFor(/* attempt= */ 5000, new Random(0))).isEqualTo(Duration.ZERO); + } + + @Test + public void delayFor_beyondTheNanosecondsALongHolds_saturates() { + Duration forever = Duration.ofSeconds(Long.MAX_VALUE); + RetryConfig config = + RetryConfig.builder().initialDelay(forever).maxDelay(forever).jitter(0).build(); + + assertThat(config.delayFor(/* attempt= */ 1, new Random(0))) + .isEqualTo(Duration.ofNanos(Long.MAX_VALUE)); + } + + private static RetryConfig withMaxAttempts(int maxAttempts) { + return RetryConfig.builder().maxAttempts(maxAttempts).build(); + } } diff --git a/core/src/test/java/com/google/adk/workflow/WorkflowResumeTest.java b/core/src/test/java/com/google/adk/workflow/WorkflowResumeTest.java index bb13a2ae1..1f9a6f027 100644 --- a/core/src/test/java/com/google/adk/workflow/WorkflowResumeTest.java +++ b/core/src/test/java/com/google/adk/workflow/WorkflowResumeTest.java @@ -20,12 +20,13 @@ import static com.google.adk.testing.WorkflowTestNodes.emitter; import static com.google.adk.testing.WorkflowTestNodes.workflow; import static com.google.adk.workflow.WorkflowResumeTesting.AGE_SCHEMA; -import static com.google.adk.workflow.WorkflowResumeTesting.LONG_RUNNING_CALL; import static com.google.adk.workflow.WorkflowResumeTesting.START; import static com.google.adk.workflow.WorkflowResumeTesting.answer; import static com.google.adk.workflow.WorkflowResumeTesting.callContent; import static com.google.adk.workflow.WorkflowResumeTesting.captureInput; import static com.google.adk.workflow.WorkflowResumeTesting.go; +import static com.google.adk.workflow.WorkflowResumeTesting.interrupt; +import static com.google.adk.workflow.WorkflowResumeTesting.longRunningAsk; import static com.google.adk.workflow.WorkflowResumeTesting.modelText; import static com.google.adk.workflow.WorkflowResumeTesting.newEvent; import static com.google.adk.workflow.WorkflowResumeTesting.outputEventAt; @@ -1360,19 +1361,6 @@ private static RunNodeConfig runId(String runId) { return RunNodeConfig.builder().runId(runId).build(); } - /** Returns a node that waits on the long-running call {@code interruptId}, as a tool does. */ - private static Node longRunningAsk(String name, String interruptId) { - return new TestNode(name, (ctx, input) -> Flowable.just(interrupt(interruptId))); - } - - /** Returns the event in which a node waits on the long-running call {@code interruptId}. */ - private static Event interrupt(String interruptId) { - return newEvent() - .content(callContent(LONG_RUNNING_CALL, interruptId)) - .longRunningToolIds(ImmutableSet.of(interruptId)) - .build(); - } - /** * Returns a node that runs {@code child} with {@code first} and then with {@code second}, whose * call completes {@code secondJoined} once it has joined a run or started its own. diff --git a/core/src/test/java/com/google/adk/workflow/WorkflowResumeTesting.java b/core/src/test/java/com/google/adk/workflow/WorkflowResumeTesting.java index a19714e10..6c9321c49 100644 --- a/core/src/test/java/com/google/adk/workflow/WorkflowResumeTesting.java +++ b/core/src/test/java/com/google/adk/workflow/WorkflowResumeTesting.java @@ -31,6 +31,7 @@ import com.google.adk.testing.TestNode; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; import com.google.genai.types.Content; import com.google.genai.types.FunctionCall; import com.google.genai.types.FunctionResponse; @@ -133,6 +134,19 @@ static Node captureInput(String name) { name, (ctx, input) -> Flowable.just(Collections.singletonMap("received", input))); } + /** Returns a node that waits on the long-running call {@code interruptId}, as a tool does. */ + static Node longRunningAsk(String name, String interruptId) { + return new TestNode(name, (ctx, input) -> Flowable.just(interrupt(interruptId))); + } + + /** Returns the event in which a node waits on the long-running call {@code interruptId}. */ + static Event interrupt(String interruptId) { + return newEvent() + .content(callContent(LONG_RUNNING_CALL, interruptId)) + .longRunningToolIds(ImmutableSet.of(interruptId)) + .build(); + } + /** A runner over an app rooted on a node, driven a turn at a time on one session. */ static final class WorkflowSession { private final Runner runner; diff --git a/core/src/test/java/com/google/adk/workflow/WorkflowRetryTest.java b/core/src/test/java/com/google/adk/workflow/WorkflowRetryTest.java new file mode 100644 index 000000000..ea4d84da2 --- /dev/null +++ b/core/src/test/java/com/google/adk/workflow/WorkflowRetryTest.java @@ -0,0 +1,771 @@ +/* + * 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.adk.testing.WorkflowTestNodes.dispatcher; +import static com.google.adk.testing.WorkflowTestNodes.throwing; +import static com.google.adk.testing.WorkflowTestNodes.workflow; +import static com.google.adk.workflow.WorkflowResumeTesting.START; +import static com.google.adk.workflow.WorkflowResumeTesting.answer; +import static com.google.adk.workflow.WorkflowResumeTesting.captureInput; +import static com.google.adk.workflow.WorkflowResumeTesting.go; +import static com.google.adk.workflow.WorkflowResumeTesting.longRunningAsk; +import static com.google.adk.workflow.WorkflowResumeTesting.modelText; +import static com.google.adk.workflow.WorkflowResumeTesting.newEvent; +import static com.google.adk.workflow.WorkflowResumeTesting.outputEventAt; +import static com.google.adk.workflow.WorkflowResumeTesting.pathsOf; +import static com.google.adk.workflow.WorkflowResumeTesting.replyInput; +import static com.google.common.collect.ImmutableList.toImmutableList; +import static com.google.common.truth.Truth.assertThat; +import static java.util.concurrent.TimeUnit.MILLISECONDS; +import static java.util.concurrent.TimeUnit.SECONDS; + +import com.google.adk.agents.Context; +import com.google.adk.agents.InvocationContext; +import com.google.adk.apps.App; +import com.google.adk.events.Event; +import com.google.adk.plugins.BasePlugin; +import com.google.adk.runner.Runner; +import com.google.adk.workflow.WorkflowResumeTesting.FailedTurn; +import com.google.adk.workflow.WorkflowResumeTesting.WorkflowSession; +import com.google.common.base.Stopwatch; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.genai.types.FinishReason; +import io.opentelemetry.context.ContextKey; +import io.opentelemetry.context.Scope; +import io.reactivex.rxjava3.core.Flowable; +import io.reactivex.rxjava3.core.Maybe; +import io.reactivex.rxjava3.core.Scheduler; +import io.reactivex.rxjava3.schedulers.Schedulers; +import java.time.Duration; +import java.util.List; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicBoolean; +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; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * Tests for the retry policy and timeout that a node's {@link NodeConfig} sets, run end to end + * through a {@link Runner}. + */ +@RunWith(JUnit4.class) +public final class WorkflowRetryTest { + + private static final RunNodeConfig NAMED_RUN = RunNodeConfig.builder().runId("n").build(); + private static final Duration SHORT_TIMEOUT = Duration.ofSeconds(1); + private static final String SCHEDULER_THREAD = "invocation-scheduler"; + + @Test + public void aFailingNode_isRetriedUntilAnAttemptSucceeds() { + CountingNode flaky = + new CountingNode( + "flaky", retrying(3), (ctx, run) -> run < 3 ? fail() : Flowable.just("recovered")); + + ImmutableList events = new WorkflowSession(workflow(new Edge(START, flaky))).turn(go()); + + assertThat(flaky.runs()).isEqualTo(3); + assertThat(errorPaths(events)).containsExactly("wf@1/flaky@1", "wf@1/flaky@1"); + assertThat(outputEventAt(events, "wf@1/flaky@1").output()).hasValue("recovered"); + } + + @Test + public void aRetriedNode_readsItsAttemptCountFromItsContext() { + List attemptCounts = new CopyOnWriteArrayList<>(); + CountingNode flaky = + new CountingNode( + "flaky", + retrying(3), + (ctx, run) -> { + attemptCounts.add(ctx.attemptCount()); + return run < 3 ? fail() : Flowable.just("recovered"); + }); + + new WorkflowSession(workflow(new Edge(START, flaky))).advance(go()); + + assertThat(attemptCounts).containsExactly(1, 2, 3).inOrder(); + } + + @Test + public void aNodeThatKeepsFailing_failsTheTurnWithItsLastFailure() { + CountingNode doomed = + new CountingNode( + "doomed", + retrying(3), + (ctx, run) -> Flowable.error(new IllegalStateException("attempt " + run))); + + FailedTurn turn = new WorkflowSession(workflow(new Edge(START, doomed))).failingTurn(go()); + + assertThat(doomed.runs()).isEqualTo(3); + assertThat(errorPaths(turn.events())).hasSize(3); + assertThat(turn.error()).hasMessageThat().isEqualTo("attempt 3"); + } + + @Test + public void aFailureOfATypeTheExceptionsDoNotList_isNotRetried() { + RetryConfig onlyTimeouts = + instantRetry(3).toBuilder().exceptions(ImmutableList.of("TimeoutException")).build(); + CountingNode flaky = + new CountingNode( + "flaky", NodeConfig.builder().retryConfig(onlyTimeouts).build(), (ctx, run) -> fail()); + + FailedTurn turn = new WorkflowSession(workflow(new Edge(START, flaky))).failingTurn(go()); + + assertThat(flaky.runs()).isEqualTo(1); + assertThat(turn.error()).isInstanceOf(IllegalStateException.class); + } + + @Test + public void aRootNode_isRetried() { + CountingNode flaky = + new CountingNode( + "flaky", retrying(2), (ctx, run) -> run == 1 ? fail() : Flowable.just("recovered")); + + ImmutableList events = new WorkflowSession(flaky).turn(go()); + + assertThat(flaky.runs()).isEqualTo(2); + assertThat(outputEventAt(events, "flaky@1").output()).hasValue("recovered"); + } + + @Test + public void aRetry_waitsOutTheBackoffDelayFirst() { + RetryConfig policy = + RetryConfig.builder().maxAttempts(2).initialDelay(Duration.ofMillis(200)).jitter(0).build(); + Stopwatch sinceFirstRun = Stopwatch.createUnstarted(); + AtomicReference gap = new AtomicReference<>(); + CountingNode flaky = + new CountingNode( + "flaky", + NodeConfig.builder().retryConfig(policy).build(), + (ctx, run) -> { + if (run == 1) { + sinceFirstRun.start(); + return fail(); + } + gap.set(sinceFirstRun.elapsed()); + return Flowable.just("recovered"); + }); + + new WorkflowSession(workflow(new Edge(START, flaky))).advance(go()); + + assertThat(gap.get()).isAtLeast(Duration.ofMillis(200)); + } + + @Test + public void aTimedAttemptAndARetry_keepTheCallersTraceContext() { + ContextKey key = ContextKey.named("caller"); + List seen = new CopyOnWriteArrayList<>(); + CountingNode timed = + new CountingNode( + "timed", + NodeConfig.builder().timeout(Duration.ofSeconds(10)).build(), + (ctx, run) -> { + seen.add("timed: " + io.opentelemetry.context.Context.current().get(key)); + return Flowable.just("in time"); + }); + CountingNode flaky = + new CountingNode( + "flaky", + retrying(2), + (ctx, run) -> { + seen.add("flaky " + run + ": " + io.opentelemetry.context.Context.current().get(key)); + return run == 1 ? fail() : Flowable.just("recovered"); + }); + + try (Scope unused = io.opentelemetry.context.Context.current().with(key, "set").makeCurrent()) { + new WorkflowSession(workflow(new Edge(START, timed), new Edge(timed, flaky))).advance(go()); + } + + assertThat(seen).containsExactly("timed: set", "flaky 1: set", "flaky 2: set").inOrder(); + } + + @Test + public void aTimeout_reportsTheFailureInTheCallersTraceContext() { + ContextKey key = ContextKey.named("caller"); + List seen = new CopyOnWriteArrayList<>(); + BasePlugin recorder = + new BasePlugin("recorder") { + @Override + public Maybe onEventCallback(InvocationContext invocationContext, Event event) { + if (event.errorCode().isPresent()) { + seen.add(String.valueOf(io.opentelemetry.context.Context.current().get(key))); + } + return Maybe.empty(); + } + }; + CountingNode slow = + new CountingNode( + "slow", NodeConfig.builder().timeout(SHORT_TIMEOUT).build(), (ctx, run) -> late()); + WorkflowSession session = + new WorkflowSession(workflow(new Edge(START, slow)), /* resumable= */ true, recorder); + + Throwable error; + try (Scope unused = io.opentelemetry.context.Context.current().with(key, "set").makeCurrent()) { + error = session.failingTurn(go()).error(); + } + + assertThat(error).isInstanceOf(TimeoutException.class); + assertThat(seen).containsExactly("set"); + } + + @Test + public void aNodeOverrunningItsTimeout_failsWithATimeoutExceptionAndIsCanceled() { + AtomicBoolean canceled = new AtomicBoolean(); + CountingNode slow = + new CountingNode( + "slow", + NodeConfig.builder().timeout(SHORT_TIMEOUT).build(), + (ctx, run) -> late().doOnCancel(() -> canceled.set(true))); + + FailedTurn turn = new WorkflowSession(workflow(new Edge(START, slow))).failingTurn(go()); + + assertThat(turn.error()).isInstanceOf(TimeoutException.class); + assertThat(turn.error()).hasMessageThat().contains("'slow'"); + assertThat(errorCodes(turn.events())).containsExactly(new FinishReason("TimeoutException")); + assertThat(canceled.get()).isTrue(); + } + + @Test + public void aTimedOutAttempt_isRetried() { + NodeConfig config = + NodeConfig.builder().retryConfig(instantRetry(2)).timeout(SHORT_TIMEOUT).build(); + CountingNode slow = + new CountingNode( + "slow", config, (ctx, run) -> run == 1 ? late() : Flowable.just("in time")); + + ImmutableList events = new WorkflowSession(workflow(new Edge(START, slow))).turn(go()); + + assertThat(slow.runs()).isEqualTo(2); + assertThat(errorCodes(events)).containsExactly(new FinishReason("TimeoutException")); + assertThat(outputEventAt(events, "wf@1/slow@1").output()).hasValue("in time"); + } + + @Test + public void aRetry_startsOnTheInvocationsScheduler() { + ExecutorService executor = + Executors.newSingleThreadExecutor(task -> new Thread(task, SCHEDULER_THREAD)); + List threads = new CopyOnWriteArrayList<>(); + // The first attempt fails on a timer's thread, which the retry must not start on. + CountingNode flaky = + new CountingNode( + "flaky", + retrying(2), + (ctx, run) -> { + threads.add(Thread.currentThread().getName()); + return run == 1 + ? Flowable.timer(10, MILLISECONDS).flatMap(unused -> fail()) + : Flowable.just("recovered"); + }); + + List events; + try { + events = + turnOn(Schedulers.from(executor), workflow(new Edge(START, flaky))) + .toList() + .blockingGet(); + } finally { + executor.shutdownNow(); + } + + assertThat(outputEventAt(events, "wf@1/flaky@1").output()).hasValue("recovered"); + assertThat(threads).hasSize(2); + assertThat(threads.get(1)).isEqualTo(SCHEDULER_THREAD); + } + + @Test + public void aTimeout_failsANodeThatHoldsTheInvocationsOnlyThread() { + ExecutorService executor = + Executors.newSingleThreadExecutor(task -> new Thread(task, SCHEDULER_THREAD)); + CountDownLatch release = new CountDownLatch(1); + AtomicBoolean returned = new AtomicBoolean(); + // The node blocks the thread it runs on, the invocation's only one, until released. + CountingNode stuck = + new CountingNode( + "stuck", + NodeConfig.builder().timeout(SHORT_TIMEOUT).build(), + (ctx, run) -> + Flowable.fromCallable( + () -> { + boolean unused = release.await(10, SECONDS); + returned.set(true); + return "late"; + })); + + AtomicReference error = new AtomicReference<>(); + boolean heldWhenTheTurnFailed; + try { + turnOn(Schedulers.from(executor), workflow(new Edge(START, stuck))) + .blockingSubscribe(event -> {}, error::set); + heldWhenTheTurnFailed = !returned.get(); + } finally { + release.countDown(); + executor.shutdownNow(); + } + + assertThat(error.get()).isInstanceOf(TimeoutException.class); + assertThat(error.get()).hasMessageThat().contains("'stuck'"); + assertThat(heldWhenTheTurnFailed).isTrue(); + } + + @Test + public void aTimeout_doesNotInterruptTheHandlingOfTheNodesEvent() { + CountingNode quick = + new CountingNode( + "quick", + NodeConfig.builder().timeout(Duration.ofMillis(300)).build(), + (ctx, run) -> Flowable.just("out")); + AtomicBoolean interrupted = new AtomicBoolean(); + AtomicReference error = new AtomicReference<>(); + + turnOn(Schedulers.io(), workflow(new Edge(START, quick))) + .doOnNext( + event -> { + if (WorkflowResumeTesting.pathOf(event).equals("wf@1/quick@1")) { + try { + // Outlasts the timeout on the thread that delivers the node's event. + Thread.sleep(1000); + } catch (InterruptedException e) { + interrupted.set(true); + Thread.currentThread().interrupt(); + } + } + }) + .blockingSubscribe(event -> {}, error::set); + + assertThat(error.get()).isInstanceOf(TimeoutException.class); + assertThat(interrupted.get()).isFalse(); + } + + @Test + public void anOuterTimeout_isNeitherBlamedOnNorRetriedByTheNodeItCancels() { + CountingNode inner = + new CountingNode( + "inner", + NodeConfig.builder() + .retryConfig(instantRetry(3)) + .timeout(Duration.ofSeconds(10)) + .build(), + (ctx, run) -> late()); + Workflow outer = + Workflow.builder() + .name("wf") + .edges(new Edge(START, inner)) + .config(NodeConfig.builder().timeout(SHORT_TIMEOUT).build()) + .build(); + + FailedTurn turn = new WorkflowSession(outer).failingTurn(go()); + + assertThat(turn.error()).hasMessageThat().contains("'wf'"); + assertThat(inner.runs()).isEqualTo(1); + assertThat(errorPaths(turn.events())).containsExactly("wf@1"); + } + + @Test + public void aWorkflowRetriedAfterANodeFailed_replaysTheNodesThatFinished() { + CountingNode first = new CountingNode("first", (ctx, run) -> Flowable.just("first_out")); + CountingNode flaky = + new CountingNode("flaky", (ctx, run) -> run == 1 ? fail() : Flowable.just("recovered")); + Workflow root = + Workflow.builder() + .name("wf") + .edges(new Edge(START, first), new Edge(first, flaky)) + .config(retrying(2)) + .build(); + + ImmutableList events = new WorkflowSession(root).turn(go()); + + assertThat(first.runs()).isEqualTo(1); + assertThat(flaky.runs()).isEqualTo(2); + assertThat(outputEventAt(events, "wf@1/flaky@1").output()).hasValue("recovered"); + } + + @Test + public void aRetriedWorkflow_rerunsANodeThatTheFailureCanceledMidRun() { + CountingNode done = new CountingNode("done", (ctx, run) -> Flowable.just("done_out")); + CountingNode stalled = stalledOnItsFirstRun("stalled"); + + ImmutableList events = + new WorkflowSession(workflow(new Edge(START, forkWithAFlakyBranch("sub", done, stalled)))) + .turn(go()); + + assertThat(done.runs()).isEqualTo(1); + assertThat(stalled.runs()).isEqualTo(2); + assertThat(outputEventAt(events, "wf@1/sub@1/join@1").output()) + .hasValue( + ImmutableMap.of("done", "done_out", "stalled", "stalled_out", "flaky", "recovered")); + } + + @Test + public void aRetriedRootWorkflow_rerunsANodeThatTheFailureCanceledMidRun() { + CountingNode done = new CountingNode("done", (ctx, run) -> Flowable.just("done_out")); + CountingNode stalled = stalledOnItsFirstRun("stalled"); + + ImmutableList events = + new WorkflowSession(forkWithAFlakyBranch("wf", done, stalled)).turn(go()); + + assertThat(done.runs()).isEqualTo(1); + assertThat(stalled.runs()).isEqualTo(2); + assertThat(outputEventAt(events, "wf@1/join@1").output()) + .hasValue( + ImmutableMap.of("done", "done_out", "stalled", "stalled_out", "flaky", "recovered")); + } + + @Test + public void aTimedOutRootNode_rerunsANamedChildThatItsTimeoutCanceledMidRun() { + CountingNode child = stalledOnItsFirstRun("child"); + CountingNode caller = + new CountingNode( + "caller", + NodeConfig.builder().retryConfig(instantRetry(2)).timeout(SHORT_TIMEOUT).build(), + (ctx, run) -> ctx.runNode(child, null, NAMED_RUN).toFlowable()); + + ImmutableList events = new WorkflowSession(caller).turn(go()); + + assertThat(child.runs()).isEqualTo(2); + assertThat(outputEventAt(events, "caller@1").output()).hasValue("child_out"); + } + + @Test + public void aRetriedWorkflow_rerunsANodeThatTookOnAChildFailureAfterReportingProgress() { + CountingNode child = + new CountingNode("child", (ctx, run) -> run == 1 ? fail() : Flowable.just("child_out")); + CountingNode caller = + new CountingNode( + "caller", + (ctx, run) -> + Flowable.just(progress()) + .concatWith(ctx.runNode(child, null).toFlowable())); + Workflow root = + Workflow.builder().name("wf").edges(new Edge(START, caller)).config(retrying(2)).build(); + + ImmutableList events = new WorkflowSession(root).turn(go()); + + assertThat(caller.runs()).isEqualTo(2); + assertThat(outputEventAt(events, "wf@1/caller@1").output()).hasValue("child_out"); + } + + @Test + public void aRetriedWorkflow_rerunsANodeThatTookOnAChildFailureAfterItsOutput() { + CountingNode child = + new CountingNode("child", (ctx, run) -> run == 1 ? fail() : Flowable.just("child_out")); + CountingNode caller = + new CountingNode( + "caller", + (ctx, run) -> + Flowable.just("caller_out_" + run) + .concatWith(ctx.runNode(child, null).ignoreElement())); + Workflow root = + Workflow.builder() + .name("wf") + .edges(new Edge(START, caller), new Edge(caller, captureInput("reply"))) + .config(retrying(2)) + .build(); + + ImmutableList events = new WorkflowSession(root).turn(go()); + + assertThat(caller.runs()).isEqualTo(2); + assertThat(replyInput(events)).isEqualTo("caller_out_2"); + } + + @Test + public void aRetriedNode_rerunsANamedChildThatTookOnAFailureAfterReportingProgress() { + CountingNode grandchild = + new CountingNode( + "grandchild", (ctx, run) -> run == 1 ? fail() : Flowable.just("grandchild_out")); + CountingNode child = + new CountingNode( + "child", + (ctx, run) -> + Flowable.just(progress()) + .concatWith(ctx.runNode(grandchild, null).toFlowable())); + CountingNode caller = + new CountingNode( + "caller", + retrying(2), + (ctx, run) -> + ctx.runNode(child, null, NAMED_RUN) + .onErrorResumeNext( + error -> Maybe.error(new IllegalStateException("child failed"))) + .toFlowable()); + + ImmutableList events = new WorkflowSession(workflow(new Edge(START, caller))).turn(go()); + + assertThat(child.runs()).isEqualTo(2); + assertThat(outputEventAt(events, "wf@1/caller@1").output()).hasValue("grandchild_out"); + } + + @Test + public void aChildWorkflowRunAgainUnderItsRunId_rerunsANodeThatTookOnAFailure() { + CountingNode grandchild = + new CountingNode( + "grandchild", (ctx, run) -> run == 1 ? fail() : Flowable.just("grandchild_out")); + CountingNode inner = + new CountingNode( + "inner", + (ctx, run) -> + Flowable.just(progress()) + .concatWith(ctx.runNode(grandchild, null).toFlowable())); + Workflow child = Workflow.builder().name("child").edges(new Edge(START, inner)).build(); + CountingNode caller = + new CountingNode( + "caller", + (ctx, run) -> + ctx.runNode(child, null, NAMED_RUN) + .onErrorResumeNext(error -> ctx.runNode(child, null, NAMED_RUN)) + .toFlowable()); + + ImmutableList events = new WorkflowSession(workflow(new Edge(START, caller))).turn(go()); + + assertThat(inner.runs()).isEqualTo(2); + assertThat(outputEventAt(events, "wf@1/caller@1").output()).hasValue("grandchild_out"); + } + + @Test + public void resume_aNodeThatSucceededOnItsRetry_replaysItsOutputWithoutRunningAgain() { + CountingNode flaky = + new CountingNode( + "flaky", retrying(2), (ctx, run) -> run == 1 ? fail() : Flowable.just("recovered")); + Node ask = longRunningAsk("ask", "approve_1"); + WorkflowSession session = + new WorkflowSession( + workflow( + new Edge(START, flaky), + new Edge(flaky, ask), + new Edge(ask, captureInput("reply")))); + session.advance(go()); + + ImmutableList events = + session.turn(answer("approve_1", ImmutableMap.of("result", "yes"))); + + assertThat(flaky.runs()).isEqualTo(2); + assertThat(outputEventAt(events, "wf@1/flaky@1").output()).hasValue("recovered"); + assertThat(replyInput(events)).isEqualTo("yes"); + } + + @Test + public void aChildWithARetryPolicy_isRetriedWithinItsCallersRun() { + CountingNode child = + new CountingNode( + "child", retrying(2), (ctx, run) -> run == 1 ? fail() : Flowable.just("child_out")); + Node caller = dispatcher("caller", (ctx, input) -> ctx.runNode(child, null).toFlowable()); + + ImmutableList events = new WorkflowSession(workflow(new Edge(START, caller))).turn(go()); + + assertThat(child.runs()).isEqualTo(2); + assertThat(outputEventAt(events, "wf@1/caller@1").output()).hasValue("child_out"); + } + + @Test + public void aNodeWhoseChildFails_isNotRetriedForThatFailure() { + Node child = throwing("child", new IllegalStateException("child failed")); + CountingNode caller = + new CountingNode( + "caller", retrying(3), (ctx, run) -> ctx.runNode(child, null).toFlowable()); + + FailedTurn turn = new WorkflowSession(workflow(new Edge(START, caller))).failingTurn(go()); + + assertThat(caller.runs()).isEqualTo(1); + assertThat(errorPaths(turn.events())).containsExactly("wf@1/caller@1/child@1"); + } + + @Test + public void aRetriedNode_reusesItsNamedChildAndRunsAnUnnamedOneAgain() { + CountingNode named = new CountingNode("named", (ctx, run) -> Flowable.just("named_out")); + CountingNode unnamed = new CountingNode("unnamed", (ctx, run) -> Flowable.just("unnamed_out")); + CountingNode caller = + new CountingNode( + "caller", + retrying(2), + (ctx, run) -> + ctx.runNode(named, null, NAMED_RUN) + .flatMap( + namedOutput -> + ctx.runNode(unnamed, null) + .map(unnamedOutput -> ImmutableList.of(namedOutput, unnamedOutput))) + .flatMap( + outputs -> + run == 1 + ? Maybe.>error( + new IllegalStateException("first attempt")) + : Maybe.just(outputs)) + .toFlowable()); + + ImmutableList events = new WorkflowSession(workflow(new Edge(START, caller))).turn(go()); + + assertThat(named.runs()).isEqualTo(1); + assertThat(unnamed.runs()).isEqualTo(2); + assertThat(pathsOf(events)).contains("wf@1/caller@1/unnamed@2"); + assertThat(outputEventAt(events, "wf@1/caller@1").output()) + .hasValue(ImmutableList.of("named_out", "unnamed_out")); + } + + @Test + public void aRetriedNode_reusesANamedChildThatFinishedWithoutAnOutput() { + CountingNode named = new CountingNode("named", (ctx, run) -> Flowable.empty()); + CountingNode caller = + new CountingNode( + "caller", + retrying(2), + (ctx, run) -> + ctx.runNode(named, null, NAMED_RUN) + .ignoreElement() + .andThen(run == 1 ? fail() : Flowable.just("done"))); + + ImmutableList events = new WorkflowSession(workflow(new Edge(START, caller))).turn(go()); + + assertThat(named.runs()).isEqualTo(1); + assertThat(outputEventAt(events, "wf@1/caller@1").output()).hasValue("done"); + } + + /** + * Returns a workflow with a retry policy that forks to {@code done}, {@code stalled} and a node + * that fails its first run, and joins the three; the failure cancels {@code stalled}. + */ + private static Workflow forkWithAFlakyBranch(String name, Node done, Node stalled) { + Node flaky = + new CountingNode("flaky", (ctx, run) -> run == 1 ? fail() : Flowable.just("recovered")); + Node join = JoinNode.builder().name("join").build(); + return Workflow.builder() + .name(name) + .edges( + new Edge(START, done), + new Edge(START, stalled), + new Edge(START, flaky), + new Edge(done, join), + new Edge(stalled, join), + new Edge(flaky, join)) + .config(retrying(2)) + .build(); + } + + /** + * Returns a node that reports progress and never finishes on its first run, so the session shows + * it started, and on later runs outputs its name followed by {@code _out}. + */ + private static CountingNode stalledOnItsFirstRun(String name) { + return new CountingNode( + name, + (ctx, run) -> + run == 1 + ? Flowable.just(progress()).concatWith(Flowable.never()) + : Flowable.just(name + "_out")); + } + + /** Returns an event that reports progress without an output, so the session shows a run. */ + private static Event progress() { + return newEvent().content(modelText("working")).build(); + } + + /** Runs one turn of an app rooted on {@code root} whose invocations use {@code scheduler}. */ + private static Flowable turnOn(Scheduler scheduler, Node root) { + Runner runner = + Runner.builder() + .app(App.builder().name("retry_app").rootNode(root).build()) + .scheduler(scheduler) + .build(); + String sessionId = + runner.sessionService().createSession("retry_app", "user").blockingGet().id(); + return runner.runAsync("user", sessionId, go()).timeout(30, SECONDS); + } + + /** Returns a retry policy of {@code maxAttempts} attempts that never waits between them. */ + private static RetryConfig instantRetry(int maxAttempts) { + return RetryConfig.builder() + .maxAttempts(maxAttempts) + .initialDelay(Duration.ZERO) + .jitter(0) + .build(); + } + + private static NodeConfig retrying(int maxAttempts) { + return NodeConfig.builder().retryConfig(instantRetry(maxAttempts)).build(); + } + + private static Flowable fail() { + return Flowable.error(new IllegalStateException("flaky")); + } + + /** Returns a run that outputs well after any timeout these tests set. */ + private static Flowable late() { + return Flowable.timer(5, SECONDS).map(unused -> "late"); + } + + private static ImmutableList errorPaths(List events) { + return events.stream() + .filter(event -> event.errorCode().isPresent()) + .map(WorkflowResumeTesting::pathOf) + .collect(toImmutableList()); + } + + private static ImmutableList errorCodes(List events) { + return events.stream().flatMap(event -> event.errorCode().stream()).collect(toImmutableList()); + } + + /** + * A node with {@code config} whose body also receives the number of the run, counting from 1, so + * a test can vary its runs and count them. It reruns on resume, so it can run child nodes. + */ + private static final class CountingNode implements Node { + private final String name; + private final NodeConfig config; + private final BiFunction> body; + private final AtomicInteger runs = new AtomicInteger(); + + CountingNode(String name, BiFunction> body) { + this(name, NodeConfig.builder().build(), body); + } + + CountingNode(String name, NodeConfig config, BiFunction> body) { + this.name = name; + this.config = config; + this.body = body; + } + + int runs() { + return runs.get(); + } + + @Override + public String name() { + return name; + } + + @Override + public NodeConfig config() { + return config; + } + + @Override + public boolean rerunOnResume() { + return true; + } + + @Override + public Flowable runNode(Context context, @Nullable Object nodeInput) { + return body.apply(context, runs.incrementAndGet()); + } + } +}