Skip to content

Commit 38e2ef4

Browse files
014-code014-code
committed
feat(core): add opt-in retries for transient LLM failures
1 parent e5aba3a commit 38e2ef4

7 files changed

Lines changed: 608 additions & 1 deletion

File tree

‎core/src/main/java/com/google/adk/agents/RetryConfig.java‎

Lines changed: 179 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,179 @@
1+
/*
2+
* Copyright 2026 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.adk.agents;
18+
19+
import com.google.common.collect.ImmutableSet;
20+
import com.google.errorprone.annotations.CanIgnoreReturnValue;
21+
import java.time.Duration;
22+
import java.util.Objects;
23+
import java.util.Set;
24+
25+
/** Configuration for retrying transient failures at the LLM provider-call boundary. */
26+
public final class RetryConfig {
27+
28+
private static final ImmutableSet<Integer> DEFAULT_RETRYABLE_STATUS_CODES =
29+
ImmutableSet.of(408, 429, 500, 502, 503, 504);
30+
31+
private final int maxAttempts;
32+
private final Duration initialBackoff;
33+
private final Duration maxBackoff;
34+
private final double multiplier;
35+
private final double jitterRatio;
36+
private final ImmutableSet<Integer> retryableStatusCodes;
37+
38+
private RetryConfig(Builder builder) {
39+
this.maxAttempts = builder.maxAttempts;
40+
this.initialBackoff = builder.initialBackoff;
41+
this.maxBackoff = builder.maxBackoff;
42+
this.multiplier = builder.multiplier;
43+
this.jitterRatio = builder.jitterRatio;
44+
this.retryableStatusCodes = ImmutableSet.copyOf(builder.retryableStatusCodes);
45+
}
46+
47+
/** Returns a retry configuration that performs exactly one provider attempt. */
48+
public static RetryConfig disabled() {
49+
return builder().build();
50+
}
51+
52+
/** Returns a builder whose defaults keep retries disabled. */
53+
public static Builder builder() {
54+
return new Builder();
55+
}
56+
57+
public int maxAttempts() {
58+
return maxAttempts;
59+
}
60+
61+
public Duration initialBackoff() {
62+
return initialBackoff;
63+
}
64+
65+
public Duration maxBackoff() {
66+
return maxBackoff;
67+
}
68+
69+
public double multiplier() {
70+
return multiplier;
71+
}
72+
73+
public double jitterRatio() {
74+
return jitterRatio;
75+
}
76+
77+
public ImmutableSet<Integer> retryableStatusCodes() {
78+
return retryableStatusCodes;
79+
}
80+
81+
@Override
82+
public boolean equals(Object other) {
83+
if (this == other) {
84+
return true;
85+
}
86+
if (!(other instanceof RetryConfig that)) {
87+
return false;
88+
}
89+
return maxAttempts == that.maxAttempts
90+
&& Double.compare(multiplier, that.multiplier) == 0
91+
&& Double.compare(jitterRatio, that.jitterRatio) == 0
92+
&& initialBackoff.equals(that.initialBackoff)
93+
&& maxBackoff.equals(that.maxBackoff)
94+
&& retryableStatusCodes.equals(that.retryableStatusCodes);
95+
}
96+
97+
@Override
98+
public int hashCode() {
99+
return Objects.hash(
100+
maxAttempts, initialBackoff, maxBackoff, multiplier, jitterRatio, retryableStatusCodes);
101+
}
102+
103+
/** Builder for {@link RetryConfig}. */
104+
public static final class Builder {
105+
private int maxAttempts = 1;
106+
private Duration initialBackoff = Duration.ofMillis(500);
107+
private Duration maxBackoff = Duration.ofSeconds(8);
108+
private double multiplier = 2.0;
109+
private double jitterRatio = 0.0;
110+
private Set<Integer> retryableStatusCodes = DEFAULT_RETRYABLE_STATUS_CODES;
111+
112+
private Builder() {}
113+
114+
@CanIgnoreReturnValue
115+
public Builder maxAttempts(int maxAttempts) {
116+
this.maxAttempts = maxAttempts;
117+
return this;
118+
}
119+
120+
@CanIgnoreReturnValue
121+
public Builder initialBackoff(Duration initialBackoff) {
122+
this.initialBackoff = Objects.requireNonNull(initialBackoff, "initialBackoff cannot be null");
123+
return this;
124+
}
125+
126+
@CanIgnoreReturnValue
127+
public Builder maxBackoff(Duration maxBackoff) {
128+
this.maxBackoff = Objects.requireNonNull(maxBackoff, "maxBackoff cannot be null");
129+
return this;
130+
}
131+
132+
@CanIgnoreReturnValue
133+
public Builder multiplier(double multiplier) {
134+
this.multiplier = multiplier;
135+
return this;
136+
}
137+
138+
@CanIgnoreReturnValue
139+
public Builder jitterRatio(double jitterRatio) {
140+
this.jitterRatio = jitterRatio;
141+
return this;
142+
}
143+
144+
@CanIgnoreReturnValue
145+
public Builder retryableStatusCodes(Set<Integer> retryableStatusCodes) {
146+
this.retryableStatusCodes =
147+
Objects.requireNonNull(retryableStatusCodes, "retryableStatusCodes cannot be null");
148+
return this;
149+
}
150+
151+
public RetryConfig build() {
152+
if (maxAttempts < 1) {
153+
throw new IllegalArgumentException("maxAttempts must be at least 1.");
154+
}
155+
if (initialBackoff.isNegative()) {
156+
throw new IllegalArgumentException("initialBackoff cannot be negative.");
157+
}
158+
if (maxBackoff.isNegative()) {
159+
throw new IllegalArgumentException("maxBackoff cannot be negative.");
160+
}
161+
if (maxBackoff.compareTo(initialBackoff) < 0) {
162+
throw new IllegalArgumentException("maxBackoff cannot be less than initialBackoff.");
163+
}
164+
if (!Double.isFinite(multiplier) || multiplier < 1.0) {
165+
throw new IllegalArgumentException("multiplier must be finite and at least 1.0.");
166+
}
167+
if (!Double.isFinite(jitterRatio) || jitterRatio < 0.0 || jitterRatio > 1.0) {
168+
throw new IllegalArgumentException("jitterRatio must be between 0.0 and 1.0.");
169+
}
170+
for (int statusCode : retryableStatusCodes) {
171+
if (statusCode < 100 || statusCode > 599) {
172+
throw new IllegalArgumentException(
173+
"retryableStatusCodes must contain valid HTTP status codes.");
174+
}
175+
}
176+
return new RetryConfig(this);
177+
}
178+
}
179+
}

‎core/src/main/java/com/google/adk/agents/RunConfig.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,8 @@ public enum ToolExecutionMode {
8585

8686
public abstract int maxLlmCalls();
8787

88+
public abstract RetryConfig retryConfig();
89+
8890
public abstract boolean autoCreateSession();
8991

9092
/**
@@ -127,6 +129,7 @@ public static Builder builder() {
127129
.streamingMode(StreamingMode.NONE)
128130
.toolExecutionMode(ToolExecutionMode.NONE)
129131
.maxLlmCalls(500)
132+
.retryConfig(RetryConfig.disabled())
130133
.autoCreateSession(false)
131134
.customMetadata(ImmutableMap.of());
132135
}
@@ -138,6 +141,7 @@ public static Builder builder(RunConfig runConfig) {
138141
.streamingMode(runConfig.streamingMode())
139142
.toolExecutionMode(runConfig.toolExecutionMode())
140143
.maxLlmCalls(runConfig.maxLlmCalls())
144+
.retryConfig(runConfig.retryConfig())
141145
.responseModalities(runConfig.responseModalities())
142146
.speechConfig(runConfig.speechConfig())
143147
.avatarConfig(runConfig.avatarConfig())
@@ -232,6 +236,9 @@ public final Builder setMaxLlmCalls(int maxLlmCalls) {
232236
@CanIgnoreReturnValue
233237
public abstract Builder maxLlmCalls(int maxLlmCalls);
234238

239+
@CanIgnoreReturnValue
240+
public abstract Builder retryConfig(RetryConfig retryConfig);
241+
235242
@Deprecated
236243
@CanIgnoreReturnValue
237244
public final Builder setAutoCreateSession(boolean autoCreateSession) {

‎core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -257,7 +257,9 @@ private Flowable<Event> callLlm(
257257
agent.resolvedModel().modelName().get());
258258
LlmRequest finalLlmRequest = llmRequestBuilder.build();
259259

260-
return llm.generateContent(
260+
return LlmRetryPolicy.execute(
261+
context,
262+
llm,
261263
finalLlmRequest,
262264
context.runConfig().streamingMode()
263265
== StreamingMode.SSE)

‎core/src/main/java/com/google/adk/flows/llmflows/LlmRetryPolicy.java‎

Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,126 @@
1+
/*
2+
* Copyright 2026 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.adk.flows.llmflows;
18+
19+
import com.google.adk.agents.InvocationContext;
20+
import com.google.adk.agents.RetryConfig;
21+
import com.google.adk.models.BaseLlm;
22+
import com.google.adk.models.LlmCallsLimitExceededException;
23+
import com.google.adk.models.LlmRequest;
24+
import com.google.adk.models.LlmResponse;
25+
import io.opentelemetry.api.trace.Span;
26+
import io.reactivex.rxjava3.core.Flowable;
27+
import java.time.Duration;
28+
import java.util.Collections;
29+
import java.util.IdentityHashMap;
30+
import java.util.Set;
31+
import java.util.concurrent.ThreadLocalRandom;
32+
import java.util.concurrent.TimeUnit;
33+
import java.util.concurrent.atomic.AtomicBoolean;
34+
import org.slf4j.Logger;
35+
import org.slf4j.LoggerFactory;
36+
37+
/** Executes retries around a single LLM provider call without replaying agent-side effects. */
38+
final class LlmRetryPolicy {
39+
private static final Logger logger = LoggerFactory.getLogger(LlmRetryPolicy.class);
40+
41+
private LlmRetryPolicy() {}
42+
43+
static Flowable<LlmResponse> execute(
44+
InvocationContext context, BaseLlm llm, LlmRequest request, boolean stream) {
45+
return executeAttempt(context, llm, request, stream, /* attempt= */ 1);
46+
}
47+
48+
private static Flowable<LlmResponse> executeAttempt(
49+
InvocationContext context, BaseLlm llm, LlmRequest request, boolean stream, int attempt) {
50+
return Flowable.defer(
51+
() -> {
52+
AtomicBoolean emittedResponse = new AtomicBoolean(false);
53+
return llm.generateContent(request, stream)
54+
.doOnNext(unused -> emittedResponse.set(true))
55+
.onErrorResumeNext(
56+
error -> {
57+
RetryConfig config = context.runConfig().retryConfig();
58+
if (emittedResponse.get()
59+
|| attempt >= config.maxAttempts()
60+
|| !isRetryable(llm, error, config)) {
61+
return Flowable.error(error);
62+
}
63+
64+
long delayMillis = retryDelay(config, attempt);
65+
int nextAttempt = attempt + 1;
66+
recordRetry(llm, context, error, nextAttempt, delayMillis);
67+
return Flowable.timer(delayMillis, TimeUnit.MILLISECONDS)
68+
.flatMapPublisher(
69+
unused -> {
70+
try {
71+
context.incrementLlmCallsCount();
72+
} catch (LlmCallsLimitExceededException e) {
73+
return Flowable.error(e);
74+
}
75+
return executeAttempt(context, llm, request, stream, nextAttempt);
76+
});
77+
});
78+
});
79+
}
80+
81+
private static boolean isRetryable(BaseLlm llm, Throwable error, RetryConfig config) {
82+
Set<Throwable> visited = Collections.newSetFromMap(new IdentityHashMap<>());
83+
for (Throwable current = error;
84+
current != null && visited.add(current);
85+
current = current.getCause()) {
86+
if (current instanceof LlmCallsLimitExceededException) {
87+
return false;
88+
}
89+
if (llm.isExceptionRetryable(current, config.retryableStatusCodes())) {
90+
return true;
91+
}
92+
}
93+
return false;
94+
}
95+
96+
private static long retryDelay(RetryConfig config, int attempt) {
97+
Duration initial = config.initialBackoff();
98+
Duration maximum = config.maxBackoff();
99+
double exponential = initial.toMillis() * Math.pow(config.multiplier(), attempt - 1);
100+
double bounded = Math.min(exponential, maximum.toMillis());
101+
if (config.jitterRatio() > 0.0 && bounded > 0.0) {
102+
double jitter =
103+
ThreadLocalRandom.current().nextDouble(-config.jitterRatio(), config.jitterRatio());
104+
bounded *= 1.0 + jitter;
105+
}
106+
return Math.max(0L, Math.round(Math.min(bounded, maximum.toMillis())));
107+
}
108+
109+
private static void recordRetry(
110+
BaseLlm llm, InvocationContext context, Throwable error, int nextAttempt, long delayMillis) {
111+
logger.warn(
112+
"Retrying LLM call for model {} and agent {} (attempt {}, delay {} ms) after {}",
113+
llm.model(),
114+
context.agent().name(),
115+
nextAttempt,
116+
delayMillis,
117+
error.getClass().getName());
118+
119+
Span span = Span.current();
120+
span.setAttribute("adk.llm.retry.model", llm.model());
121+
span.setAttribute("adk.llm.retry.agent", context.agent().name());
122+
span.setAttribute("adk.llm.retry.attempt", nextAttempt);
123+
span.setAttribute("adk.llm.retry.delay_ms", delayMillis);
124+
span.setAttribute("adk.llm.retry.error_type", error.getClass().getName());
125+
}
126+
}

0 commit comments

Comments
 (0)