From 38e2ef472d1f074aea88f89d384cba241a85128a Mon Sep 17 00:00:00 2001 From: 014-code <161032298+014-code@users.noreply.github.com> Date: Sat, 8 Aug 2026 21:23:20 +0800 Subject: [PATCH] feat(core): add opt-in retries for transient LLM failures --- .../com/google/adk/agents/RetryConfig.java | 179 ++++++++++++++ .../java/com/google/adk/agents/RunConfig.java | 7 + .../adk/flows/llmflows/BaseLlmFlow.java | 4 +- .../adk/flows/llmflows/LlmRetryPolicy.java | 126 ++++++++++ .../java/com/google/adk/models/BaseLlm.java | 38 +++ .../com/google/adk/agents/RunConfigTest.java | 37 +++ .../adk/flows/llmflows/BaseLlmFlowTest.java | 218 ++++++++++++++++++ 7 files changed, 608 insertions(+), 1 deletion(-) create mode 100644 core/src/main/java/com/google/adk/agents/RetryConfig.java create mode 100644 core/src/main/java/com/google/adk/flows/llmflows/LlmRetryPolicy.java diff --git a/core/src/main/java/com/google/adk/agents/RetryConfig.java b/core/src/main/java/com/google/adk/agents/RetryConfig.java new file mode 100644 index 000000000..638ddd6e8 --- /dev/null +++ b/core/src/main/java/com/google/adk/agents/RetryConfig.java @@ -0,0 +1,179 @@ +/* + * 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.agents; + +import com.google.common.collect.ImmutableSet; +import com.google.errorprone.annotations.CanIgnoreReturnValue; +import java.time.Duration; +import java.util.Objects; +import java.util.Set; + +/** Configuration for retrying transient failures at the LLM provider-call boundary. */ +public final class RetryConfig { + + private static final ImmutableSet DEFAULT_RETRYABLE_STATUS_CODES = + ImmutableSet.of(408, 429, 500, 502, 503, 504); + + private final int maxAttempts; + private final Duration initialBackoff; + private final Duration maxBackoff; + private final double multiplier; + private final double jitterRatio; + private final ImmutableSet retryableStatusCodes; + + private RetryConfig(Builder builder) { + this.maxAttempts = builder.maxAttempts; + this.initialBackoff = builder.initialBackoff; + this.maxBackoff = builder.maxBackoff; + this.multiplier = builder.multiplier; + this.jitterRatio = builder.jitterRatio; + this.retryableStatusCodes = ImmutableSet.copyOf(builder.retryableStatusCodes); + } + + /** Returns a retry configuration that performs exactly one provider attempt. */ + public static RetryConfig disabled() { + return builder().build(); + } + + /** Returns a builder whose defaults keep retries disabled. */ + public static Builder builder() { + return new Builder(); + } + + public int maxAttempts() { + return maxAttempts; + } + + public Duration initialBackoff() { + return initialBackoff; + } + + public Duration maxBackoff() { + return maxBackoff; + } + + public double multiplier() { + return multiplier; + } + + public double jitterRatio() { + return jitterRatio; + } + + public ImmutableSet retryableStatusCodes() { + return retryableStatusCodes; + } + + @Override + public boolean equals(Object other) { + if (this == other) { + return true; + } + if (!(other instanceof RetryConfig that)) { + return false; + } + return maxAttempts == that.maxAttempts + && Double.compare(multiplier, that.multiplier) == 0 + && Double.compare(jitterRatio, that.jitterRatio) == 0 + && initialBackoff.equals(that.initialBackoff) + && maxBackoff.equals(that.maxBackoff) + && retryableStatusCodes.equals(that.retryableStatusCodes); + } + + @Override + public int hashCode() { + return Objects.hash( + maxAttempts, initialBackoff, maxBackoff, multiplier, jitterRatio, retryableStatusCodes); + } + + /** Builder for {@link RetryConfig}. */ + public static final class Builder { + private int maxAttempts = 1; + private Duration initialBackoff = Duration.ofMillis(500); + private Duration maxBackoff = Duration.ofSeconds(8); + private double multiplier = 2.0; + private double jitterRatio = 0.0; + private Set retryableStatusCodes = DEFAULT_RETRYABLE_STATUS_CODES; + + private Builder() {} + + @CanIgnoreReturnValue + public Builder maxAttempts(int maxAttempts) { + this.maxAttempts = maxAttempts; + return this; + } + + @CanIgnoreReturnValue + public Builder initialBackoff(Duration initialBackoff) { + this.initialBackoff = Objects.requireNonNull(initialBackoff, "initialBackoff cannot be null"); + return this; + } + + @CanIgnoreReturnValue + public Builder maxBackoff(Duration maxBackoff) { + this.maxBackoff = Objects.requireNonNull(maxBackoff, "maxBackoff cannot be null"); + return this; + } + + @CanIgnoreReturnValue + public Builder multiplier(double multiplier) { + this.multiplier = multiplier; + return this; + } + + @CanIgnoreReturnValue + public Builder jitterRatio(double jitterRatio) { + this.jitterRatio = jitterRatio; + return this; + } + + @CanIgnoreReturnValue + public Builder retryableStatusCodes(Set retryableStatusCodes) { + this.retryableStatusCodes = + Objects.requireNonNull(retryableStatusCodes, "retryableStatusCodes cannot be null"); + return this; + } + + public RetryConfig build() { + if (maxAttempts < 1) { + throw new IllegalArgumentException("maxAttempts must be at least 1."); + } + if (initialBackoff.isNegative()) { + throw new IllegalArgumentException("initialBackoff cannot be negative."); + } + if (maxBackoff.isNegative()) { + throw new IllegalArgumentException("maxBackoff cannot be negative."); + } + if (maxBackoff.compareTo(initialBackoff) < 0) { + throw new IllegalArgumentException("maxBackoff cannot be less than initialBackoff."); + } + if (!Double.isFinite(multiplier) || multiplier < 1.0) { + throw new IllegalArgumentException("multiplier must be finite and at least 1.0."); + } + if (!Double.isFinite(jitterRatio) || jitterRatio < 0.0 || jitterRatio > 1.0) { + throw new IllegalArgumentException("jitterRatio must be between 0.0 and 1.0."); + } + for (int statusCode : retryableStatusCodes) { + if (statusCode < 100 || statusCode > 599) { + throw new IllegalArgumentException( + "retryableStatusCodes must contain valid HTTP status codes."); + } + } + return new RetryConfig(this); + } + } +} diff --git a/core/src/main/java/com/google/adk/agents/RunConfig.java b/core/src/main/java/com/google/adk/agents/RunConfig.java index bd20b6183..411e4633d 100644 --- a/core/src/main/java/com/google/adk/agents/RunConfig.java +++ b/core/src/main/java/com/google/adk/agents/RunConfig.java @@ -85,6 +85,8 @@ public enum ToolExecutionMode { public abstract int maxLlmCalls(); + public abstract RetryConfig retryConfig(); + public abstract boolean autoCreateSession(); /** @@ -127,6 +129,7 @@ public static Builder builder() { .streamingMode(StreamingMode.NONE) .toolExecutionMode(ToolExecutionMode.NONE) .maxLlmCalls(500) + .retryConfig(RetryConfig.disabled()) .autoCreateSession(false) .customMetadata(ImmutableMap.of()); } @@ -138,6 +141,7 @@ public static Builder builder(RunConfig runConfig) { .streamingMode(runConfig.streamingMode()) .toolExecutionMode(runConfig.toolExecutionMode()) .maxLlmCalls(runConfig.maxLlmCalls()) + .retryConfig(runConfig.retryConfig()) .responseModalities(runConfig.responseModalities()) .speechConfig(runConfig.speechConfig()) .avatarConfig(runConfig.avatarConfig()) @@ -232,6 +236,9 @@ public final Builder setMaxLlmCalls(int maxLlmCalls) { @CanIgnoreReturnValue public abstract Builder maxLlmCalls(int maxLlmCalls); + @CanIgnoreReturnValue + public abstract Builder retryConfig(RetryConfig retryConfig); + @Deprecated @CanIgnoreReturnValue public final Builder setAutoCreateSession(boolean autoCreateSession) { diff --git a/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java b/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java index 91cc225f2..9aaf69c1a 100644 --- a/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java +++ b/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java @@ -257,7 +257,9 @@ private Flowable callLlm( agent.resolvedModel().modelName().get()); LlmRequest finalLlmRequest = llmRequestBuilder.build(); - return llm.generateContent( + return LlmRetryPolicy.execute( + context, + llm, finalLlmRequest, context.runConfig().streamingMode() == StreamingMode.SSE) diff --git a/core/src/main/java/com/google/adk/flows/llmflows/LlmRetryPolicy.java b/core/src/main/java/com/google/adk/flows/llmflows/LlmRetryPolicy.java new file mode 100644 index 000000000..13c99e57a --- /dev/null +++ b/core/src/main/java/com/google/adk/flows/llmflows/LlmRetryPolicy.java @@ -0,0 +1,126 @@ +/* + * 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.flows.llmflows; + +import com.google.adk.agents.InvocationContext; +import com.google.adk.agents.RetryConfig; +import com.google.adk.models.BaseLlm; +import com.google.adk.models.LlmCallsLimitExceededException; +import com.google.adk.models.LlmRequest; +import com.google.adk.models.LlmResponse; +import io.opentelemetry.api.trace.Span; +import io.reactivex.rxjava3.core.Flowable; +import java.time.Duration; +import java.util.Collections; +import java.util.IdentityHashMap; +import java.util.Set; +import java.util.concurrent.ThreadLocalRandom; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** Executes retries around a single LLM provider call without replaying agent-side effects. */ +final class LlmRetryPolicy { + private static final Logger logger = LoggerFactory.getLogger(LlmRetryPolicy.class); + + private LlmRetryPolicy() {} + + static Flowable execute( + InvocationContext context, BaseLlm llm, LlmRequest request, boolean stream) { + return executeAttempt(context, llm, request, stream, /* attempt= */ 1); + } + + private static Flowable executeAttempt( + InvocationContext context, BaseLlm llm, LlmRequest request, boolean stream, int attempt) { + return Flowable.defer( + () -> { + AtomicBoolean emittedResponse = new AtomicBoolean(false); + return llm.generateContent(request, stream) + .doOnNext(unused -> emittedResponse.set(true)) + .onErrorResumeNext( + error -> { + RetryConfig config = context.runConfig().retryConfig(); + if (emittedResponse.get() + || attempt >= config.maxAttempts() + || !isRetryable(llm, error, config)) { + return Flowable.error(error); + } + + long delayMillis = retryDelay(config, attempt); + int nextAttempt = attempt + 1; + recordRetry(llm, context, error, nextAttempt, delayMillis); + return Flowable.timer(delayMillis, TimeUnit.MILLISECONDS) + .flatMapPublisher( + unused -> { + try { + context.incrementLlmCallsCount(); + } catch (LlmCallsLimitExceededException e) { + return Flowable.error(e); + } + return executeAttempt(context, llm, request, stream, nextAttempt); + }); + }); + }); + } + + private static boolean isRetryable(BaseLlm llm, Throwable error, RetryConfig config) { + Set visited = Collections.newSetFromMap(new IdentityHashMap<>()); + for (Throwable current = error; + current != null && visited.add(current); + current = current.getCause()) { + if (current instanceof LlmCallsLimitExceededException) { + return false; + } + if (llm.isExceptionRetryable(current, config.retryableStatusCodes())) { + return true; + } + } + return false; + } + + private static long retryDelay(RetryConfig config, int attempt) { + Duration initial = config.initialBackoff(); + Duration maximum = config.maxBackoff(); + double exponential = initial.toMillis() * Math.pow(config.multiplier(), attempt - 1); + double bounded = Math.min(exponential, maximum.toMillis()); + if (config.jitterRatio() > 0.0 && bounded > 0.0) { + double jitter = + ThreadLocalRandom.current().nextDouble(-config.jitterRatio(), config.jitterRatio()); + bounded *= 1.0 + jitter; + } + return Math.max(0L, Math.round(Math.min(bounded, maximum.toMillis()))); + } + + private static void recordRetry( + BaseLlm llm, InvocationContext context, Throwable error, int nextAttempt, long delayMillis) { + logger.warn( + "Retrying LLM call for model {} and agent {} (attempt {}, delay {} ms) after {}", + llm.model(), + context.agent().name(), + nextAttempt, + delayMillis, + error.getClass().getName()); + + Span span = Span.current(); + span.setAttribute("adk.llm.retry.model", llm.model()); + span.setAttribute("adk.llm.retry.agent", context.agent().name()); + span.setAttribute("adk.llm.retry.attempt", nextAttempt); + span.setAttribute("adk.llm.retry.delay_ms", delayMillis); + span.setAttribute("adk.llm.retry.error_type", error.getClass().getName()); + } +} diff --git a/core/src/main/java/com/google/adk/models/BaseLlm.java b/core/src/main/java/com/google/adk/models/BaseLlm.java index f57dfe4af..0b116fa79 100644 --- a/core/src/main/java/com/google/adk/models/BaseLlm.java +++ b/core/src/main/java/com/google/adk/models/BaseLlm.java @@ -16,7 +16,12 @@ package com.google.adk.models; +import com.google.genai.errors.GenAiIOException; import io.reactivex.rxjava3.core.Flowable; +import java.io.IOException; +import java.lang.reflect.Method; +import java.util.Set; +import java.util.concurrent.TimeoutException; /** * Abstract base class for Large Language Models (LLMs). @@ -52,6 +57,39 @@ public String model() { */ public abstract Flowable generateContent(LlmRequest llmRequest, boolean stream); + /** + * Returns whether a provider failure is safe to retry. + * + *

Implementations may override this method for provider-specific exception types. The default + * implementation recognizes common network failures and exceptions exposing an HTTP status code + * through {@code code()}, {@code statusCode()}, or {@code getStatusCode()}. + */ + public boolean isExceptionRetryable(Throwable exception, Set retryableStatusCodes) { + if (exception instanceof IOException + || exception instanceof TimeoutException + || exception instanceof GenAiIOException) { + return true; + } + + Integer statusCode = extractStatusCode(exception); + return statusCode != null && retryableStatusCodes.contains(statusCode); + } + + private static Integer extractStatusCode(Throwable exception) { + for (String methodName : new String[] {"code", "statusCode", "getStatusCode"}) { + try { + Method method = exception.getClass().getMethod(methodName); + Object value = method.invoke(exception); + if (value instanceof Number number) { + return number.intValue(); + } + } catch (ReflectiveOperationException | SecurityException ignored) { + // Try the next provider convention. + } + } + return null; + } + /** Creates a live connection to the LLM. */ public abstract BaseLlmConnection connect(LlmRequest llmRequest); } diff --git a/core/src/test/java/com/google/adk/agents/RunConfigTest.java b/core/src/test/java/com/google/adk/agents/RunConfigTest.java index fc6b9083f..7575401c6 100644 --- a/core/src/test/java/com/google/adk/agents/RunConfigTest.java +++ b/core/src/test/java/com/google/adk/agents/RunConfigTest.java @@ -25,7 +25,9 @@ import com.google.genai.types.CustomizedAvatar; import com.google.genai.types.Modality; import com.google.genai.types.SpeechConfig; +import java.time.Duration; import java.util.Optional; +import java.util.Set; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -57,6 +59,7 @@ public void testBuilderWithVariousValues() { assertThat(runConfig.outputAudioTranscription()).isEqualTo(audioTranscriptionConfig); assertThat(runConfig.inputAudioTranscription()).isEqualTo(audioTranscriptionConfig); assertThat(runConfig.maxLlmCalls()).isEqualTo(10); + assertThat(runConfig.retryConfig().maxAttempts()).isEqualTo(1); } @Test @@ -71,6 +74,7 @@ public void testBuilderDefaults() { assertThat(runConfig.outputAudioTranscription()).isNull(); assertThat(runConfig.inputAudioTranscription()).isNull(); assertThat(runConfig.maxLlmCalls()).isEqualTo(500); + assertThat(runConfig.retryConfig()).isEqualTo(RetryConfig.disabled()); assertThat(runConfig.autoCreateSession()).isFalse(); assertThat(runConfig.groupFunctionResponsesInHistoryOverride()).isEmpty(); assertThat(runConfig.groupFunctionResponsesInHistory()).isFalse(); @@ -160,6 +164,39 @@ public void testMaxLlmCalls_integerMaxValue_throwsIllegalArgumentException() { () -> RunConfig.builder().setMaxLlmCalls(Integer.MAX_VALUE).build()); } + @Test + public void retryConfig_customValues_areAppliedAndCopied() { + RetryConfig retryConfig = + RetryConfig.builder() + .maxAttempts(3) + .initialBackoff(Duration.ofMillis(100)) + .maxBackoff(Duration.ofSeconds(2)) + .multiplier(1.5) + .jitterRatio(0.2) + .retryableStatusCodes(Set.of(429, 503)) + .build(); + + RunConfig runConfig = RunConfig.builder().retryConfig(retryConfig).build(); + RunConfig copied = RunConfig.builder(runConfig).build(); + + assertThat(copied.retryConfig()).isEqualTo(retryConfig); + } + + @Test + public void retryConfig_invalidValues_throwIllegalArgumentException() { + assertThrows( + IllegalArgumentException.class, () -> RetryConfig.builder().maxAttempts(0).build()); + assertThrows( + IllegalArgumentException.class, + () -> + RetryConfig.builder() + .initialBackoff(Duration.ofSeconds(2)) + .maxBackoff(Duration.ofSeconds(1)) + .build()); + assertThrows( + IllegalArgumentException.class, () -> RetryConfig.builder().jitterRatio(1.1).build()); + } + @Test public void testAvatarConfig_withName() { AvatarConfig avatarConfig = AvatarConfig.builder().avatarName("test_avatar").build(); diff --git a/core/src/test/java/com/google/adk/flows/llmflows/BaseLlmFlowTest.java b/core/src/test/java/com/google/adk/flows/llmflows/BaseLlmFlowTest.java index 1761871e6..663ffcbf3 100644 --- a/core/src/test/java/com/google/adk/flows/llmflows/BaseLlmFlowTest.java +++ b/core/src/test/java/com/google/adk/flows/llmflows/BaseLlmFlowTest.java @@ -31,10 +31,13 @@ import com.google.adk.agents.InvocationContext; import com.google.adk.agents.LlmAgent; import com.google.adk.agents.ReadonlyContext; +import com.google.adk.agents.RetryConfig; import com.google.adk.agents.RunConfig; import com.google.adk.events.Event; import com.google.adk.flows.llmflows.RequestProcessor.RequestProcessingResult; import com.google.adk.flows.llmflows.ResponseProcessor.ResponseProcessingResult; +import com.google.adk.models.BaseLlm; +import com.google.adk.models.BaseLlmConnection; import com.google.adk.models.LlmRequest; import com.google.adk.models.LlmResponse; import com.google.adk.testing.TestLlm; @@ -59,11 +62,14 @@ import io.reactivex.rxjava3.core.Maybe; import io.reactivex.rxjava3.core.Single; import io.reactivex.rxjava3.schedulers.Schedulers; +import java.io.IOException; import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.Set; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Function; +import java.util.function.IntFunction; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -89,6 +95,182 @@ public void run_singleTextResponse_returnsSingleEvent() { assertThat(event.usageMetadata()).isEmpty(); } + @Test + public void run_retryConfig_retriesRetryableFailure() { + Content content = Content.fromParts(Part.fromText("LLM response")); + RetryTestLlm testLlm = + new RetryTestLlm( + attempt -> + attempt == 1 + ? Flowable.error(new RuntimeException(new IOException("temporary"))) + : Flowable.just(createLlmResponse(content))); + RunConfig runConfig = + RunConfig.builder() + .retryConfig( + RetryConfig.builder() + .maxAttempts(3) + .initialBackoff(java.time.Duration.ZERO) + .maxBackoff(java.time.Duration.ZERO) + .build()) + .build(); + InvocationContext invocationContext = + createInvocationContext(createTestAgent(testLlm), runConfig); + + List events = + createBaseLlmFlowWithoutProcessors().run(invocationContext).toList().blockingGet(); + + assertThat(events).hasSize(1); + assertThat(events.get(0).content()).hasValue(content); + assertThat(testLlm.attempts()).isEqualTo(2); + } + + @Test + public void run_retryConfig_retriesConfiguredStatusCode() { + Content content = Content.fromParts(Part.fromText("LLM response")); + RetryTestLlm testLlm = + new RetryTestLlm( + attempt -> + attempt == 1 + ? Flowable.error(new StatusCodeException(429)) + : Flowable.just(createLlmResponse(content))); + RunConfig runConfig = + RunConfig.builder() + .retryConfig( + RetryConfig.builder() + .maxAttempts(2) + .initialBackoff(java.time.Duration.ZERO) + .maxBackoff(java.time.Duration.ZERO) + .retryableStatusCodes(Set.of(429)) + .build()) + .build(); + InvocationContext invocationContext = + createInvocationContext(createTestAgent(testLlm), runConfig); + + List events = + createBaseLlmFlowWithoutProcessors().run(invocationContext).toList().blockingGet(); + + assertThat(events).hasSize(1); + assertThat(testLlm.attempts()).isEqualTo(2); + } + + @Test + public void run_retryConfig_disabledByDefault_doesNotRetry() { + RetryTestLlm testLlm = + new RetryTestLlm(attempt -> Flowable.error(new IOException("temporary"))); + InvocationContext invocationContext = createInvocationContext(createTestAgent(testLlm)); + + createBaseLlmFlowWithoutProcessors() + .run(invocationContext) + .test() + .assertError(IOException.class); + + assertThat(testLlm.attempts()).isEqualTo(1); + } + + @Test + public void run_retryConfig_doesNotRetryNonRetryableFailure() { + RetryTestLlm testLlm = + new RetryTestLlm( + attempt -> Flowable.error(new IllegalArgumentException("invalid request"))); + RunConfig runConfig = + RunConfig.builder() + .retryConfig( + RetryConfig.builder() + .maxAttempts(3) + .initialBackoff(java.time.Duration.ZERO) + .maxBackoff(java.time.Duration.ZERO) + .build()) + .build(); + InvocationContext invocationContext = + createInvocationContext(createTestAgent(testLlm), runConfig); + + createBaseLlmFlowWithoutProcessors() + .run(invocationContext) + .test() + .assertError(IllegalArgumentException.class); + + assertThat(testLlm.attempts()).isEqualTo(1); + } + + @Test + public void run_retryConfig_stopsAfterMaxAttempts() { + RetryTestLlm testLlm = + new RetryTestLlm(attempt -> Flowable.error(new IOException("temporary"))); + RunConfig runConfig = + RunConfig.builder() + .retryConfig( + RetryConfig.builder() + .maxAttempts(3) + .initialBackoff(java.time.Duration.ZERO) + .maxBackoff(java.time.Duration.ZERO) + .build()) + .build(); + InvocationContext invocationContext = + createInvocationContext(createTestAgent(testLlm), runConfig); + + createBaseLlmFlowWithoutProcessors() + .run(invocationContext) + .test() + .assertError(IOException.class); + + assertThat(testLlm.attempts()).isEqualTo(3); + } + + @Test + public void run_retryConfig_countsEachProviderAttemptAgainstLlmCallLimit() { + RetryTestLlm testLlm = + new RetryTestLlm(attempt -> Flowable.error(new IOException("temporary"))); + RunConfig runConfig = + RunConfig.builder() + .maxLlmCalls(1) + .retryConfig( + RetryConfig.builder() + .maxAttempts(3) + .initialBackoff(java.time.Duration.ZERO) + .maxBackoff(java.time.Duration.ZERO) + .build()) + .build(); + InvocationContext invocationContext = + createInvocationContext(createTestAgent(testLlm), runConfig); + + createBaseLlmFlowWithoutProcessors() + .run(invocationContext) + .test() + .assertError(com.google.adk.models.LlmCallsLimitExceededException.class); + + assertThat(testLlm.attempts()).isEqualTo(1); + } + + @Test + public void run_retryConfig_doesNotRetryAfterStreamingResponseWasEmitted() { + Content partial = Content.fromParts(Part.fromText("partial")); + RetryTestLlm testLlm = + new RetryTestLlm( + attempt -> + Flowable.concat( + Flowable.just(LlmResponse.builder().content(partial).partial(true).build()), + Flowable.error(new IOException("stream interrupted")))); + RunConfig runConfig = + RunConfig.builder() + .streamingMode(RunConfig.StreamingMode.SSE) + .retryConfig( + RetryConfig.builder() + .maxAttempts(3) + .initialBackoff(java.time.Duration.ZERO) + .maxBackoff(java.time.Duration.ZERO) + .build()) + .build(); + InvocationContext invocationContext = + createInvocationContext(createTestAgent(testLlm), runConfig); + + createBaseLlmFlowWithoutProcessors() + .run(invocationContext) + .test() + .assertError(IOException.class); + + assertThat(testLlm.attempts()).isEqualTo(1); + } + @Test public void run_singleTextResponse_withMetadata_returnsSingleEventWithMetadata() { Content content = Content.fromParts(Part.fromText("LLM response")); @@ -682,6 +864,42 @@ public Single> runAsync(Map args, ToolContex } } + private static final class RetryTestLlm extends BaseLlm { + private final AtomicInteger attempts = new AtomicInteger(); + private final IntFunction> responseFactory; + + RetryTestLlm(IntFunction> responseFactory) { + super("retry-test-llm"); + this.responseFactory = responseFactory; + } + + @Override + public Flowable generateContent(LlmRequest llmRequest, boolean stream) { + return responseFactory.apply(attempts.incrementAndGet()); + } + + @Override + public BaseLlmConnection connect(LlmRequest llmRequest) { + throw new UnsupportedOperationException(); + } + + int attempts() { + return attempts.get(); + } + } + + public static final class StatusCodeException extends RuntimeException { + private final int code; + + StatusCodeException(int code) { + this.code = code; + } + + public int code() { + return code; + } + } + @Test public void run_contextPropagation() { ContextKey testKey = ContextKey.named("test-key");