This is an automated email from the ASF dual-hosted git repository.
hengyuss pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/shenyu.git
The following commit(s) were added to refs/heads/master by this push:
new 17e68d7f4a feat(ai-proxy): harden configurable retry and fallback
policy (#7668)
17e68d7f4a is described below
commit 17e68d7f4a183037b193fd0b4ca24645f5dac15a
Author: hengyuss <[email protected]>
AuthorDate: Tue Oct 20 14:08:03 2026 +0800
feat(ai-proxy): harden configurable retry and fallback policy (#7668)
* refactor(ai-proxy): classify retry and fallback failures
* refactor(ai-proxy): centralize retry and fallback decisions
* refactor(ai-proxy): add tests
* feat(ai-proxy): configure retry policy
---
.../src/main/resources/application.yml | 7 +
.../ai/proxy/enhanced/service/AiProxyEngine.java | 107 +++++++---
.../proxy/enhanced/service/AiProxyRetryPolicy.java | 97 +++++++++
.../enhanced/service/AiProxyRetryProperties.java | 134 ++++++++++++
.../ai/proxy/enhanced/AiProxyPluginTest.java | 10 +-
.../service/AiProxyEngineOrchestrationTest.java | 224 ++++++++++++++++++++-
.../proxy/enhanced/service/AiProxyEngineTest.java | 3 +-
.../enhanced/service/AiProxyRetryPolicyTest.java | 107 ++++++++++
.../service/AiProxyRetryPropertiesTest.java | 89 ++++++++
.../ai/proxy/AiProxyPluginConfiguration.java | 19 +-
.../ai/proxy/AiProxyPluginConfigurationTest.java | 40 ++++
11 files changed, 787 insertions(+), 50 deletions(-)
diff --git a/shenyu-bootstrap/src/main/resources/application.yml
b/shenyu-bootstrap/src/main/resources/application.yml
index 0bb10820dd..bd6896e6cc 100644
--- a/shenyu-bootstrap/src/main/resources/application.yml
+++ b/shenyu-bootstrap/src/main/resources/application.yml
@@ -330,6 +330,13 @@ shenyu:
# plugins:
# ai:
# proxy:
+# retry:
+# # Maximum retries after the initial non-streaming request. Default:
3.
+# max-retries: 3
+# # Maximum retries before a streaming response emits its first
chunk. Default: 1.
+# streaming-max-retries: 1
+# # Initial backoff for non-streaming retries. Subsequent delays
double. Default: 1s.
+# initial-backoff: 1s
# transport:
# # Maximum time allowed to establish a connection to the AI upstream.
# # Default: 30s.
diff --git
a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/main/java/org/apache/shenyu/plugin/ai/proxy/enhanced/service/AiProxyEngine.java
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/main/java/org/apache/shenyu/plugin/ai/proxy/enhanced/service/AiProxyEngine.java
index e146623e06..088fb9156e 100644
---
a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/main/java/org/apache/shenyu/plugin/ai/proxy/enhanced/service/AiProxyEngine.java
+++
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/main/java/org/apache/shenyu/plugin/ai/proxy/enhanced/service/AiProxyEngine.java
@@ -50,24 +50,38 @@ public final class AiProxyEngine {
private static final ObjectMapper OBJECT_MAPPER = new
ObjectMapper().enable(DeserializationFeature.FAIL_ON_TRAILING_TOKENS);
+ private static final AiProxyRetryPolicy RETRY_POLICY = new
AiProxyRetryPolicy();
+
private final AiProxyProtocolFactory aiProxyProtocolFactory;
private final AiProxyProviderFactory aiProxyProviderFactory;
private final AiProxyTransport aiProxyTransport;
- private final Duration retryBackoff;
-
- public AiProxyEngine(final AiProxyProtocolFactory protocols, final
AiProxyProviderFactory providers, final AiProxyTransport transport) {
- this(protocols, providers, transport, Duration.ofSeconds(1));
+ private final AiProxyRetryProperties retryProperties;
+
+ public AiProxyEngine(
+ final AiProxyProtocolFactory protocols,
+ final AiProxyProviderFactory providers,
+ final AiProxyTransport transport) {
+ this(
+ protocols,
+ providers,
+ transport,
+ new AiProxyRetryProperties());
}
public AiProxyEngine(final AiProxyProtocolFactory protocols, final
AiProxyProviderFactory providers,
- final AiProxyTransport transport, final Duration
retryBackoff) {
+ final AiProxyTransport transport, final
AiProxyRetryProperties retryProperties) {
this.aiProxyProtocolFactory = protocols;
this.aiProxyProviderFactory = providers;
this.aiProxyTransport = transport;
- this.retryBackoff = retryBackoff;
+ this.retryProperties = Objects.requireNonNull(retryProperties,
"retryProperties must not be null");
+ }
+
+ public AiProxyEngine(final AiProxyProtocolFactory protocols, final
AiProxyProviderFactory providers,
+ final AiProxyTransport transport, final Duration
retryBackoff) {
+ this(protocols, providers, transport, new AiProxyRetryProperties(3, 1,
retryBackoff));
}
/**
@@ -92,7 +106,8 @@ public final class AiProxyEngine {
copy.setProtocol(primary.protocol().name());
return copy;
});
- return attempt(requestBody, primary, fallback, primary.stream() ?
1 : 3, responseConsumer, new AtomicBoolean());
+ final int maxRetries =
retryProperties.maxRetries(primary.stream());
+ return attempt(requestBody, primary, fallback, maxRetries,
maxRetries, responseConsumer, new AtomicBoolean());
});
}
@@ -108,48 +123,78 @@ public final class AiProxyEngine {
}
private Mono<Void> attempt(final String requestBody, final
PreparedExchange prepared,
- final Optional<AiCommonConfig> fallback, final
int retries,
+ final Optional<AiCommonConfig> fallback, final
int maxRetries, final int retriesRemaining,
final Function<AiUpstreamResponse, Mono<Void>>
consumer, final AtomicBoolean emitted) {
return Mono.defer(() -> aiProxyTransport.execute(prepared.request()))
.flatMap(response -> {
int status = response.getStatusCode();
- if (status >= 400 && ((retryableStatus(status) && retries
> 0) || fallback.isPresent())) {
- // Release this response without accumulating the
error body before another attempt.
+ if (status >= 400) {
+ final AiProxyRetryPolicy.Decision decision =
+ RETRY_POLICY.decideStatus(status,
retriesRemaining, fallback.isPresent());
+ if (decision ==
AiProxyRetryPolicy.Decision.FORWARD_RESPONSE) {
+ return consumeResponse(response, consumer,
emitted);
+ }
return response.getBody().then(Mono.error(new
UpstreamStatusFailure(status)));
}
AiUpstreamResponse adapted =
prepared.request().adaptResponse(prepared.provider().adaptResponse(response));
AiUpstreamResponse encoded =
prepared.protocol().encodeResponse(adapted, prepared.stream());
- return encoded.getBody().switchOnFirst((signal, body) -> {
- if (signal.isOnError()) {
- return
Flux.error(Objects.requireNonNull(signal.getThrowable()));
- }
- AiUpstreamResponse ready = new
AiUpstreamResponse(encoded.getStatusCode(), encoded.getHeaders(),
- body.doOnNext(bytes -> emitted.set(true)));
- return Mono.defer(() -> consumer.apply(ready)).flux();
- }).then();
+ return consumeResponse(encoded, consumer, emitted);
})
.onErrorResume(error -> {
- boolean transportFailure = error instanceof
AiTransportException
- || error instanceof AiProxyException
- && ((AiProxyException) error).getKind() ==
AiProxyException.Kind.TRANSPORT;
- boolean statusFailure = error instanceof
UpstreamStatusFailure;
- if (emitted.get() || !(transportFailure || statusFailure))
{
+ final AiProxyRetryPolicy.Decision decision;
+ if (emitted.get()) {
+ return Mono.error(error);
+ }
+ if (error instanceof AiTransportException
transportException) {
+ decision = RETRY_POLICY.decideTransport(
+ transportException.getKind(),
+ retriesRemaining,
+ fallback.isPresent());
+ } else if (error instanceof UpstreamStatusFailure
statusFailure) {
+ decision = RETRY_POLICY.decideStatus(
+ statusFailure.status,
+ retriesRemaining,
+ fallback.isPresent());
+ } else {
return Mono.error(error);
}
- boolean retryable = transportFailure ||
retryableStatus(((UpstreamStatusFailure) error).status);
- if (retryable && retries > 0) {
- Duration delay = prepared.stream() ? Duration.ZERO :
retryBackoff.multipliedBy(1L << (3 - retries));
- return Mono.delay(delay).then(attempt(requestBody,
prepared, fallback, retries - 1, consumer, emitted));
+ if (decision == AiProxyRetryPolicy.Decision.RETRY) {
+ final int retryIndex = maxRetries - retriesRemaining;
+ final Duration delay =
retryProperties.retryDelay(prepared.stream(), retryIndex);
+ return Mono.delay(delay).then(attempt(
+ requestBody, prepared, fallback, maxRetries,
retriesRemaining - 1, consumer, emitted));
}
- if (fallback.isPresent()) {
- return Mono.defer(() -> attempt(requestBody,
prepare(requestBody, fallback.get()), Optional.empty(), 0, consumer, emitted));
+ if (decision == AiProxyRetryPolicy.Decision.FALLBACK) {
+ return Mono.defer(() -> attempt(
+ requestBody, prepare(requestBody,
fallback.get()), Optional.empty(), 0, 0, consumer, emitted));
}
return Mono.error(error);
});
}
- private boolean retryableStatus(final int status) {
- return status == 429 || (status >= 500 && status < 600);
+ private Mono<Void> consumeResponse(
+ final AiUpstreamResponse response,
+ final Function<AiUpstreamResponse, Mono<Void>> consumer,
+ final AtomicBoolean emitted) {
+ return response.getBody()
+ .switchOnFirst((signal, body) -> {
+ if (signal.isOnError()) {
+ return Flux.error(
+ Objects.requireNonNull(
+ signal.getThrowable()));
+ }
+
+ final AiUpstreamResponse ready =
+ new AiUpstreamResponse(
+ response.getStatusCode(),
+ response.getHeaders(),
+ body.doOnNext(
+ bytes -> emitted.set(true)));
+
+ return Mono.defer(
+ () -> consumer.apply(ready)).flux();
+ })
+ .then();
}
private PreparedExchange prepare(final String requestBody, final
AiCommonConfig config) {
diff --git
a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/main/java/org/apache/shenyu/plugin/ai/proxy/enhanced/service/AiProxyRetryPolicy.java
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/main/java/org/apache/shenyu/plugin/ai/proxy/enhanced/service/AiProxyRetryPolicy.java
new file mode 100644
index 0000000000..cde94f1387
--- /dev/null
+++
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/main/java/org/apache/shenyu/plugin/ai/proxy/enhanced/service/AiProxyRetryPolicy.java
@@ -0,0 +1,97 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.shenyu.plugin.ai.proxy.enhanced.service;
+
+import
org.apache.shenyu.plugin.ai.proxy.enhanced.transport.AiTransportException;
+
+/**
+ * Decides whether an AI upstream failure should be retried, switched to
fallback,
+ * forwarded to the client, or propagated.
+ */
+final class AiProxyRetryPolicy {
+
+ Decision decideStatus(
+ final int status,
+ final int retriesRemaining,
+ final boolean fallbackAvailable) {
+ if (!retryableStatus(status)) {
+ return Decision.FORWARD_RESPONSE;
+ }
+ if (retriesRemaining > 0) {
+ return Decision.RETRY;
+ }
+ if (fallbackAvailable) {
+ return Decision.FALLBACK;
+ }
+ return Decision.FORWARD_RESPONSE;
+ }
+
+ Decision decideTransport(
+ final AiTransportException.Kind kind,
+ final int retriesRemaining,
+ final boolean fallbackAvailable) {
+ if (retriesRemaining > 0 && retryableTransportFailure(kind)) {
+ return Decision.RETRY;
+ }
+ if (fallbackAvailable && fallbackEligibleTransportFailure(kind)) {
+ return Decision.FALLBACK;
+ }
+ return Decision.PROPAGATE_ERROR;
+ }
+
+ private boolean retryableStatus(final int status) {
+ return status == 429 || status >= 500 && status < 600;
+ }
+
+ private boolean retryableTransportFailure(final AiTransportException.Kind
kind) {
+ return switch (kind) {
+ case CONNECT_TIMEOUT,
+ RESPONSE_HEADER_TIMEOUT,
+ STREAM_IDLE_TIMEOUT,
+ CONNECTION_RESET,
+ IO_FAILURE -> true;
+ case STREAM_DURATION_TIMEOUT,
+ RESPONSE_TOO_LARGE,
+ DNS_FAILURE,
+ TLS_FAILURE,
+ UNKNOWN -> false;
+ };
+ }
+
+ private boolean fallbackEligibleTransportFailure(final
AiTransportException.Kind kind) {
+ return switch (kind) {
+ case CONNECT_TIMEOUT,
+ RESPONSE_HEADER_TIMEOUT,
+ RESPONSE_TOO_LARGE,
+ CONNECTION_RESET,
+ IO_FAILURE -> true;
+ case STREAM_IDLE_TIMEOUT,
+ STREAM_DURATION_TIMEOUT,
+ DNS_FAILURE,
+ TLS_FAILURE,
+ UNKNOWN -> false;
+ };
+ }
+
+ enum Decision {
+ RETRY,
+ FALLBACK,
+ FORWARD_RESPONSE,
+ PROPAGATE_ERROR
+ }
+}
diff --git
a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/main/java/org/apache/shenyu/plugin/ai/proxy/enhanced/service/AiProxyRetryProperties.java
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/main/java/org/apache/shenyu/plugin/ai/proxy/enhanced/service/AiProxyRetryProperties.java
new file mode 100644
index 0000000000..5f1cdc7b6d
--- /dev/null
+++
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/main/java/org/apache/shenyu/plugin/ai/proxy/enhanced/service/AiProxyRetryProperties.java
@@ -0,0 +1,134 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.shenyu.plugin.ai.proxy.enhanced.service;
+
+import java.time.Duration;
+import java.util.Objects;
+
+/**
+ * Retry properties and backoff calculation for AI upstream requests.
+ */
+public final class AiProxyRetryProperties {
+
+ private static final int MAX_EXPONENTIAL_RETRIES = 30;
+
+ private int maxRetries = 3;
+
+ private int streamingMaxRetries = 1;
+
+ private Duration initialBackoff = Duration.ofSeconds(1);
+
+ /**
+ * Create retry properties with defaults.
+ */
+ public AiProxyRetryProperties() {
+ }
+
+ /**
+ * Create retry properties.
+ *
+ * @param maxRetries maximum retries for non-streaming requests
+ * @param streamingMaxRetries maximum retries before a streaming response
emits data
+ * @param initialBackoff initial backoff for non-streaming retries
+ */
+ public AiProxyRetryProperties(final int maxRetries, final int
streamingMaxRetries, final Duration initialBackoff) {
+ setMaxRetries(maxRetries);
+ setStreamingMaxRetries(streamingMaxRetries);
+ setInitialBackoff(initialBackoff);
+ }
+
+ /**
+ * Get maximum retries for non-streaming requests.
+ *
+ * @return maximum retries
+ */
+ public int getMaxRetries() {
+ return maxRetries;
+ }
+
+ /**
+ * Set maximum retries for non-streaming requests.
+ *
+ * @param maxRetries maximum retries
+ */
+ public void setMaxRetries(final int maxRetries) {
+ this.maxRetries = requireRetryCount(maxRetries, "maxRetries");
+ }
+
+ /**
+ * Get maximum retries before a streaming response emits data.
+ *
+ * @return maximum streaming retries
+ */
+ public int getStreamingMaxRetries() {
+ return streamingMaxRetries;
+ }
+
+ /**
+ * Set maximum retries before a streaming response emits data.
+ *
+ * @param streamingMaxRetries maximum streaming retries
+ */
+ public void setStreamingMaxRetries(final int streamingMaxRetries) {
+ this.streamingMaxRetries = requireRetryCount(streamingMaxRetries,
"streamingMaxRetries");
+ }
+
+ /**
+ * Get initial non-streaming retry backoff.
+ *
+ * @return initial retry backoff
+ */
+ public Duration getInitialBackoff() {
+ return initialBackoff;
+ }
+
+ /**
+ * Set initial non-streaming retry backoff.
+ *
+ * @param initialBackoff initial retry backoff
+ */
+ public void setInitialBackoff(final Duration initialBackoff) {
+ if (Objects.isNull(initialBackoff) || initialBackoff.isNegative()) {
+ throw new IllegalArgumentException("initialBackoff must not be
null or negative");
+ }
+ this.initialBackoff = initialBackoff;
+ }
+
+ int maxRetries(final boolean stream) {
+ return stream ? streamingMaxRetries : maxRetries;
+ }
+
+ Duration retryDelay(final boolean stream, final int retryIndex) {
+ if (retryIndex < 0 || retryIndex >= MAX_EXPONENTIAL_RETRIES) {
+ throw new IllegalArgumentException(
+ "retryIndex must be between 0 and " +
(MAX_EXPONENTIAL_RETRIES - 1));
+ }
+ if (stream) {
+ return Duration.ZERO;
+ }
+ return initialBackoff.multipliedBy(1L << retryIndex);
+ }
+
+ private int requireRetryCount(final int value, final String name) {
+ if (value < 0 || value > MAX_EXPONENTIAL_RETRIES) {
+ throw new IllegalArgumentException(
+ name + " must be between 0 and " +
MAX_EXPONENTIAL_RETRIES);
+ }
+ return value;
+ }
+}
diff --git
a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/test/java/org/apache/shenyu/plugin/ai/proxy/enhanced/AiProxyPluginTest.java
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/test/java/org/apache/shenyu/plugin/ai/proxy/enhanced/AiProxyPluginTest.java
index a18ff12a23..4431de85c7 100644
---
a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/test/java/org/apache/shenyu/plugin/ai/proxy/enhanced/AiProxyPluginTest.java
+++
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/test/java/org/apache/shenyu/plugin/ai/proxy/enhanced/AiProxyPluginTest.java
@@ -146,11 +146,11 @@ class AiProxyPluginTest {
@Test
void testExecuteWithAdminFallback() {
handle.setFallbackConfig(fallback("admin"));
- AiProxyPlugin plugin = plugin(request ->
Mono.just(request.getUri().getHost().equals("main.example") ? response(400,
"bad") : response(202, "admin")));
+ AiProxyPlugin plugin = plugin(request ->
Mono.just(request.getUri().getHost().equals("main.example") ? response(500,
"bad") : response(202, "admin")));
StepVerifier.create(execute(plugin)).verifyComplete();
- assertEquals(2, requests.size());
- assertEquals("admin.example", requests.get(1).getUri().getHost());
- assertEquals(List.of("Bearer admin-key"),
requests.get(1).getHeaders().get("Authorization"));
+ assertEquals(5, requests.size());
+ assertEquals("admin.example", requests.get(4).getUri().getHost());
+ assertEquals(List.of("Bearer admin-key"),
requests.get(4).getHeaders().get("Authorization"));
assertEquals(HttpStatus.ACCEPTED,
exchange.getResponse().getStatusCode());
}
@@ -308,6 +308,6 @@ class AiProxyPluginTest {
}
private AiTransportException failure() {
- return new AiTransportException(AiTransportException.Kind.UNKNOWN,
"fixture failure", null);
+ return new AiTransportException(AiTransportException.Kind.IO_FAILURE,
"fixture failure", null);
}
}
diff --git
a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/test/java/org/apache/shenyu/plugin/ai/proxy/enhanced/service/AiProxyEngineOrchestrationTest.java
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/test/java/org/apache/shenyu/plugin/ai/proxy/enhanced/service/AiProxyEngineOrchestrationTest.java
index b6746a82ca..ee485ee542 100644
---
a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/test/java/org/apache/shenyu/plugin/ai/proxy/enhanced/service/AiProxyEngineOrchestrationTest.java
+++
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/test/java/org/apache/shenyu/plugin/ai/proxy/enhanced/service/AiProxyEngineOrchestrationTest.java
@@ -32,7 +32,11 @@ import
org.apache.shenyu.plugin.ai.proxy.enhanced.provider.AiProxyProviderFactor
import org.apache.shenyu.plugin.ai.proxy.enhanced.provider.OpenAi;
import org.apache.shenyu.plugin.ai.proxy.enhanced.provider.OpenAiCompatible;
import org.apache.shenyu.plugin.ai.proxy.enhanced.transport.AiProxyTransport;
+import
org.apache.shenyu.plugin.ai.proxy.enhanced.transport.AiTransportException;
import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
+import org.junit.jupiter.params.provider.ValueSource;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.test.StepVerifier;
@@ -49,6 +53,7 @@ import java.util.concurrent.atomic.AtomicInteger;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertArrayEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
@@ -121,6 +126,54 @@ class AiProxyEngineOrchestrationTest {
assertEquals(List.of("upstream unavailable"), received);
}
+ @Test
+ void configuredRetryCountsControlPrimaryAttempts() {
+ AtomicInteger nonStreamingCalls = new AtomicInteger();
+ AiProxyEngine nonStreamingEngine = engine(request -> {
+ nonStreamingCalls.incrementAndGet();
+ return Mono.just(response(503, Flux.just(bytes("unavailable"))));
+ }, new AiProxyRetryProperties(2, 0, Duration.ZERO));
+
+ StepVerifier.create(nonStreamingEngine.execute(BODY, config(false,
"main"), Optional.empty(),
+ response -> response.getBody().then())).verifyComplete();
+ assertEquals(3, nonStreamingCalls.get());
+
+ AtomicInteger streamingCalls = new AtomicInteger();
+ AiProxyEngine streamingEngine = engine(request -> {
+ streamingCalls.incrementAndGet();
+ return Mono.error(networkFailure());
+ }, new AiProxyRetryProperties(2, 0, Duration.ZERO));
+
+ StepVerifier.create(streamingEngine.execute(BODY, config(true,
"main"), Optional.empty(),
+ response ->
response.getBody().then())).expectError(AiTransportException.class).verify();
+ assertEquals(1, streamingCalls.get());
+ }
+
+ @Test
+ void nonStreamingRetriesUseExponentialBackoff() {
+ AtomicInteger calls = new AtomicInteger();
+ AiProxyEngine engine = engine(request -> {
+ int attempt = calls.incrementAndGet();
+ return Mono.just(attempt < 4
+ ? response(503, Flux.just(bytes("unavailable")))
+ : response(200, Flux.just(bytes("ok"))));
+ }, new AiProxyRetryProperties(3, 1, Duration.ofSeconds(1)));
+
+ StepVerifier.withVirtualTime(() -> engine.execute(BODY, config(false,
"main"), Optional.empty(),
+ response -> response.getBody().then()))
+ .expectSubscription()
+ .then(() -> assertEquals(1, calls.get()))
+ .thenAwait(Duration.ofMillis(999))
+ .then(() -> assertEquals(1, calls.get()))
+ .thenAwait(Duration.ofMillis(1))
+ .then(() -> assertEquals(2, calls.get()))
+ .thenAwait(Duration.ofSeconds(2))
+ .then(() -> assertEquals(3, calls.get()))
+ .thenAwait(Duration.ofSeconds(4))
+ .then(() -> assertEquals(4, calls.get()))
+ .verifyComplete();
+ }
+
@Test
void
streamingRetriesOnceAndRebuildsFallbackUsingItsProviderConfiguration() {
List<AiUpstreamRequest> requests = new ArrayList<>();
@@ -150,18 +203,48 @@ class AiProxyEngineOrchestrationTest {
assertEquals(List.of("data: {\"vendor\":true}\n\n", "data:
[DONE]\n\n"), received);
}
- @Test
- void clientErrorsAreForwardedWithoutRetryAndMayUseConfiguredFallback() {
+ @ParameterizedTest
+ @ValueSource(ints = {400, 401, 403, 404, 408, 422})
+ void clientErrorsAreForwardedWithoutRetryOrFallback(final int status) {
AtomicInteger calls = new AtomicInteger();
+ List<String> received = new ArrayList<>();
AiProxyEngine engine = engine(request -> {
calls.incrementAndGet();
- return Mono.just(response(401, Flux.just(bytes("unauthorized"))));
+ return Mono.just(response(status, Flux.just(bytes("client
error"))));
});
- StepVerifier.create(engine.execute(BODY, config(false, "main"),
Optional.empty(), response -> response.getBody().then())).verifyComplete();
+
+ StepVerifier.create(engine.execute(BODY, config(false, "main"),
Optional.of(config(false, "fallback")), response -> {
+ assertEquals(status, response.getStatusCode());
+ assertEquals(List.of("kept"),
response.getHeaders().get("X-Upstream"));
+ return consume(response, received);
+ })).verifyComplete();
+
assertEquals(1, calls.get());
- calls.set(0);
- StepVerifier.create(engine.execute(BODY, config(false, "main"),
Optional.of(config(false, "fallback")), response ->
response.getBody().then())).verifyComplete();
- assertEquals(2, calls.get());
+ assertEquals(List.of("client error"), received);
+ }
+
+ @ParameterizedTest
+ @ValueSource(ints = {429, 500, 502, 503, 504, 599})
+ void retryableStatusesRetryBeforeUsingFallback(final int status) {
+ AtomicInteger primaryCalls = new AtomicInteger();
+ AtomicInteger fallbackCalls = new AtomicInteger();
+ List<String> received = new ArrayList<>();
+ AiProxyEngine engine = engine(request -> {
+ String requestBody = new String(request.getBody(),
StandardCharsets.UTF_8);
+ if (requestBody.contains("\"model\":\"fallback\"")) {
+ fallbackCalls.incrementAndGet();
+ return Mono.just(response(200, Flux.just(bytes("fallback
response"))));
+ }
+ primaryCalls.incrementAndGet();
+ return Mono.just(response(status, Flux.just(bytes("retryable
error"))));
+ });
+
+ StepVerifier.create(engine.execute(BODY, config(false, "main"),
Optional.of(config(false, "fallback")),
+ response -> consume(response, received))).verifyComplete();
+
+ assertEquals(4, primaryCalls.get());
+ assertEquals(1, fallbackCalls.get());
+ assertEquals(List.of("fallback response"), received);
}
@Test
@@ -175,6 +258,112 @@ class AiProxyEngineOrchestrationTest {
assertEquals(List.of("data: {\"ok\":true}\n\n", "data: [DONE]\n\n"),
received);
}
+ @ParameterizedTest
+ @CsvSource({
+ "CONNECT_TIMEOUT, 2",
+ "RESPONSE_HEADER_TIMEOUT, 2",
+ "STREAM_IDLE_TIMEOUT, 2",
+ "STREAM_DURATION_TIMEOUT, 1",
+ "RESPONSE_TOO_LARGE, 1",
+ "DNS_FAILURE, 1",
+ "TLS_FAILURE, 1",
+ "CONNECTION_RESET, 2",
+ "IO_FAILURE, 2",
+ "UNKNOWN, 1"
+ })
+ void transportFailuresFollowTheRetryWhitelist(final
AiTransportException.Kind kind, final int expectedCalls) {
+ AtomicInteger calls = new AtomicInteger();
+ AiProxyEngine engine = engine(request -> {
+ calls.incrementAndGet();
+ return Mono.error(transportFailure(kind));
+ });
+
+ StepVerifier.create(engine.execute(BODY, config(true, "main"),
Optional.empty(), response -> response.getBody().then()))
+ .expectErrorSatisfies(error -> assertEquals(kind,
+ assertInstanceOf(AiTransportException.class,
error).getKind()))
+ .verify();
+
+ assertEquals(expectedCalls, calls.get());
+ }
+
+ @ParameterizedTest
+ @CsvSource({
+ "CONNECT_TIMEOUT, 2, true",
+ "RESPONSE_HEADER_TIMEOUT, 2, true",
+ "STREAM_IDLE_TIMEOUT, 2, false",
+ "STREAM_DURATION_TIMEOUT, 1, false",
+ "RESPONSE_TOO_LARGE, 1, true",
+ "DNS_FAILURE, 1, false",
+ "TLS_FAILURE, 1, false",
+ "CONNECTION_RESET, 2, true",
+ "IO_FAILURE, 2, true",
+ "UNKNOWN, 1, false"
+ })
+ void transportFailuresFollowTheFallbackWhitelist(final
AiTransportException.Kind kind,
+ final int
expectedPrimaryCalls,
+ final boolean
fallbackExpected) {
+ AtomicInteger primaryCalls = new AtomicInteger();
+ AtomicInteger fallbackCalls = new AtomicInteger();
+ AiProxyEngine engine = engine(request -> {
+ String requestBody = new String(request.getBody(),
StandardCharsets.UTF_8);
+ if (requestBody.contains("\"model\":\"fallback\"")) {
+ fallbackCalls.incrementAndGet();
+ return Mono.just(response(200, Flux.just(bytes("fallback
response"))));
+ }
+ primaryCalls.incrementAndGet();
+ return Mono.error(transportFailure(kind));
+ });
+
+ Mono<Void> execution = engine.execute(BODY, config(true, "main"),
Optional.of(config(true, "fallback")),
+ response -> response.getBody().then());
+ if (fallbackExpected) {
+ StepVerifier.create(execution).verifyComplete();
+ } else {
+ StepVerifier.create(execution)
+ .expectErrorSatisfies(error -> assertEquals(kind,
+ assertInstanceOf(AiTransportException.class,
error).getKind()))
+ .verify();
+ }
+
+ assertEquals(expectedPrimaryCalls, primaryCalls.get());
+ assertEquals(fallbackExpected ? 1 : 0, fallbackCalls.get());
+ }
+
+ @Test
+ void responseTooLargeDoesNotRetryTheSameProvider() {
+ AtomicInteger calls = new AtomicInteger();
+ AiProxyEngine engine = engine(request -> {
+ calls.incrementAndGet();
+ return Mono.error(responseTooLargeFailure());
+ });
+
+ StepVerifier.create(engine.execute(BODY, config(false, "main"),
Optional.empty(), response -> response.getBody().then()))
+ .expectErrorSatisfies(error ->
assertEquals(AiTransportException.Kind.RESPONSE_TOO_LARGE,
+ assertInstanceOf(AiTransportException.class,
error).getKind()))
+ .verify();
+ assertEquals(1, calls.get());
+ }
+
+ @Test
+ void responseTooLargeCanUseConfiguredFallback() {
+ AtomicInteger calls = new AtomicInteger();
+ AtomicBoolean fallbackUsed = new AtomicBoolean();
+ AiProxyEngine engine = engine(request -> {
+ calls.incrementAndGet();
+ String requestBody = new String(request.getBody(),
StandardCharsets.UTF_8);
+ if (!requestBody.contains("\"model\":\"fallback\"")) {
+ return Mono.error(responseTooLargeFailure());
+ }
+ fallbackUsed.set(true);
+ return Mono.just(response(200, Flux.just(bytes("fallback
response"))));
+ });
+
+ StepVerifier.create(engine.execute(BODY, config(false, "main"),
Optional.of(config(false, "fallback")),
+ response -> response.getBody().then())).verifyComplete();
+ assertEquals(2, calls.get());
+ assertTrue(fallbackUsed.get());
+ }
+
@Test
void neverRetriesOrFallsBackAfterBodyEmission() {
AtomicInteger calls = new AtomicInteger();
@@ -184,7 +373,7 @@ class AiProxyEngineOrchestrationTest {
return Mono.just(response(200, Flux.just(bytes("data:
{\"first\":true}\n\n")).concatWith(Flux.error(networkFailure()))));
});
StepVerifier.create(engine.execute(BODY, config(true, "main"),
Optional.of(config(true, "fallback")), response -> consume(response, received)))
- .expectError(AiProxyException.class).verify();
+ .expectError(AiTransportException.class).verify();
assertEquals(1, calls.get());
assertEquals(List.of("data: {\"first\":true}\n\n"), received);
}
@@ -197,7 +386,7 @@ class AiProxyEngineOrchestrationTest {
return Mono.error(networkFailure());
});
StepVerifier.create(engine.execute(BODY, config(true, "main"),
Optional.of(config(true, "fallback")), response -> response.getBody().then()))
- .expectError(AiProxyException.class).verify();
+ .expectError(AiTransportException.class).verify();
assertEquals(3, calls.get());
}
@@ -283,6 +472,11 @@ class AiProxyEngineOrchestrationTest {
new AiProxyProviderFactory(List.of(new OpenAi(), new
OpenAiCompatible("DeepSeek"))), transport, Duration.ZERO);
}
+ private AiProxyEngine engine(final AiProxyTransport transport, final
AiProxyRetryProperties retryProperties) {
+ return new AiProxyEngine(new AiProxyProtocolFactory(List.of(new
OpenAiChat())),
+ new AiProxyProviderFactory(List.of(new OpenAi(), new
OpenAiCompatible("DeepSeek"))), transport, retryProperties);
+ }
+
private AiCommonConfig config(final boolean stream, final String model) {
AiCommonConfig config = new AiCommonConfig();
config.setProvider("OpenAI");
@@ -305,7 +499,15 @@ class AiProxyEngineOrchestrationTest {
return value.getBytes(StandardCharsets.UTF_8);
}
- private AiProxyException networkFailure() {
- return new AiProxyException(AiProxyException.Kind.TRANSPORT, "test
failure", null);
+ private AiTransportException networkFailure() {
+ return transportFailure(AiTransportException.Kind.IO_FAILURE);
+ }
+
+ private AiTransportException responseTooLargeFailure() {
+ return transportFailure(AiTransportException.Kind.RESPONSE_TOO_LARGE);
+ }
+
+ private AiTransportException transportFailure(final
AiTransportException.Kind kind) {
+ return new AiTransportException(kind, "test transport failure: " +
kind, null);
}
}
diff --git
a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/test/java/org/apache/shenyu/plugin/ai/proxy/enhanced/service/AiProxyEngineTest.java
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/test/java/org/apache/shenyu/plugin/ai/proxy/enhanced/service/AiProxyEngineTest.java
index 7efece125c..cdcc7c502c 100644
---
a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/test/java/org/apache/shenyu/plugin/ai/proxy/enhanced/service/AiProxyEngineTest.java
+++
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/test/java/org/apache/shenyu/plugin/ai/proxy/enhanced/service/AiProxyEngineTest.java
@@ -33,6 +33,7 @@ import
org.apache.shenyu.plugin.ai.proxy.enhanced.provider.AiProxyProvider;
import
org.apache.shenyu.plugin.ai.proxy.enhanced.provider.AiProxyProviderFactory;
import org.apache.shenyu.plugin.ai.proxy.enhanced.provider.OpenAi;
import org.apache.shenyu.plugin.ai.proxy.enhanced.transport.AiProxyTransport;
+import
org.apache.shenyu.plugin.ai.proxy.enhanced.transport.AiTransportException;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
@@ -212,7 +213,7 @@ class AiProxyEngineTest {
engine = new AiProxyEngine(new AiProxyProtocolFactory(List.of(new
OpenAiChat())),
new AiProxyProviderFactory(List.of(new OpenAi(), fallback)),
request -> {
transports.incrementAndGet();
- return Mono.error(new
AiProxyException(AiProxyException.Kind.TRANSPORT, "unavailable", null));
+ return Mono.error(new
AiTransportException(AiTransportException.Kind.IO_FAILURE, "unavailable",
null));
}, Duration.ZERO);
config.setStream(true);
final AiCommonConfig fallbackConfig = new AiCommonConfig(config);
diff --git
a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/test/java/org/apache/shenyu/plugin/ai/proxy/enhanced/service/AiProxyRetryPolicyTest.java
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/test/java/org/apache/shenyu/plugin/ai/proxy/enhanced/service/AiProxyRetryPolicyTest.java
new file mode 100644
index 0000000000..9933af48ec
--- /dev/null
+++
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/test/java/org/apache/shenyu/plugin/ai/proxy/enhanced/service/AiProxyRetryPolicyTest.java
@@ -0,0 +1,107 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.shenyu.plugin.ai.proxy.enhanced.service;
+
+import
org.apache.shenyu.plugin.ai.proxy.enhanced.transport.AiTransportException;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+/**
+ * Test cases for {@link AiProxyRetryPolicy}.
+ */
+class AiProxyRetryPolicyTest {
+
+ private final AiProxyRetryPolicy policy = new AiProxyRetryPolicy();
+
+ @ParameterizedTest
+ @CsvSource({
+ "400, 3, true, FORWARD_RESPONSE",
+ "401, 3, true, FORWARD_RESPONSE",
+ "408, 3, true, FORWARD_RESPONSE",
+ "422, 3, true, FORWARD_RESPONSE",
+ "429, 1, true, RETRY",
+ "429, 0, true, FALLBACK",
+ "429, 0, false, FORWARD_RESPONSE",
+ "500, 1, true, RETRY",
+ "503, 0, true, FALLBACK",
+ "599, 0, false, FORWARD_RESPONSE",
+ "600, 3, true, FORWARD_RESPONSE"
+ })
+ void shouldDecideStatusFailure(final int status,
+ final int retriesRemaining,
+ final boolean fallbackAvailable,
+ final AiProxyRetryPolicy.Decision expected)
{
+ assertEquals(expected, policy.decideStatus(status, retriesRemaining,
fallbackAvailable));
+ }
+
+ @ParameterizedTest
+ @CsvSource({
+ "CONNECT_TIMEOUT, RETRY",
+ "RESPONSE_HEADER_TIMEOUT, RETRY",
+ "STREAM_IDLE_TIMEOUT, RETRY",
+ "STREAM_DURATION_TIMEOUT, PROPAGATE_ERROR",
+ "RESPONSE_TOO_LARGE, FALLBACK",
+ "DNS_FAILURE, PROPAGATE_ERROR",
+ "TLS_FAILURE, PROPAGATE_ERROR",
+ "CONNECTION_RESET, RETRY",
+ "IO_FAILURE, RETRY",
+ "UNKNOWN, PROPAGATE_ERROR"
+ })
+ void shouldRetryBeforeUsingFallback(final AiTransportException.Kind kind,
+ final AiProxyRetryPolicy.Decision
expected) {
+ assertEquals(expected, policy.decideTransport(kind, 1, true));
+ }
+
+ @ParameterizedTest
+ @CsvSource({
+ "CONNECT_TIMEOUT, FALLBACK",
+ "RESPONSE_HEADER_TIMEOUT, FALLBACK",
+ "STREAM_IDLE_TIMEOUT, PROPAGATE_ERROR",
+ "STREAM_DURATION_TIMEOUT, PROPAGATE_ERROR",
+ "RESPONSE_TOO_LARGE, FALLBACK",
+ "DNS_FAILURE, PROPAGATE_ERROR",
+ "TLS_FAILURE, PROPAGATE_ERROR",
+ "CONNECTION_RESET, FALLBACK",
+ "IO_FAILURE, FALLBACK",
+ "UNKNOWN, PROPAGATE_ERROR"
+ })
+ void shouldApplyFallbackWhitelistAfterRetriesAreExhausted(final
AiTransportException.Kind kind,
+ final
AiProxyRetryPolicy.Decision expected) {
+ assertEquals(expected, policy.decideTransport(kind, 0, true));
+ }
+
+ @ParameterizedTest
+ @CsvSource({
+ "CONNECT_TIMEOUT",
+ "RESPONSE_HEADER_TIMEOUT",
+ "STREAM_IDLE_TIMEOUT",
+ "STREAM_DURATION_TIMEOUT",
+ "RESPONSE_TOO_LARGE",
+ "DNS_FAILURE",
+ "TLS_FAILURE",
+ "CONNECTION_RESET",
+ "IO_FAILURE",
+ "UNKNOWN"
+ })
+ void shouldPropagateTransportFailureWithoutRetriesOrFallback(final
AiTransportException.Kind kind) {
+ assertEquals(AiProxyRetryPolicy.Decision.PROPAGATE_ERROR,
+ policy.decideTransport(kind, 0, false));
+ }
+}
diff --git
a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/test/java/org/apache/shenyu/plugin/ai/proxy/enhanced/service/AiProxyRetryPropertiesTest.java
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/test/java/org/apache/shenyu/plugin/ai/proxy/enhanced/service/AiProxyRetryPropertiesTest.java
new file mode 100644
index 0000000000..c85ea36f92
--- /dev/null
+++
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/test/java/org/apache/shenyu/plugin/ai/proxy/enhanced/service/AiProxyRetryPropertiesTest.java
@@ -0,0 +1,89 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.shenyu.plugin.ai.proxy.enhanced.service;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
+
+import java.time.Duration;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+/**
+ * Test cases for {@link AiProxyRetryProperties}.
+ */
+class AiProxyRetryPropertiesTest {
+
+ @Test
+ void shouldProvideCompatibleDefaults() {
+ final AiProxyRetryProperties properties = new AiProxyRetryProperties();
+
+ assertEquals(3, properties.maxRetries(false));
+ assertEquals(1, properties.maxRetries(true));
+ assertEquals(Duration.ofSeconds(1), properties.getInitialBackoff());
+ }
+
+ @ParameterizedTest
+ @CsvSource({
+ "0, 1000",
+ "1, 2000",
+ "2, 4000"
+ })
+ void shouldCalculateExponentialBackoff(final int retryIndex, final long
expectedMillis) {
+ final AiProxyRetryProperties properties = new
AiProxyRetryProperties(3, 1, Duration.ofSeconds(1));
+
+ assertEquals(Duration.ofMillis(expectedMillis),
properties.retryDelay(false, retryIndex));
+ }
+
+ @ParameterizedTest
+ @CsvSource({"0", "1", "29"})
+ void shouldRetryStreamingRequestsImmediately(final int retryIndex) {
+ final AiProxyRetryProperties properties = new
AiProxyRetryProperties(3, 1, Duration.ofSeconds(1));
+
+ assertEquals(Duration.ZERO, properties.retryDelay(true, retryIndex));
+ }
+
+ @ParameterizedTest
+ @CsvSource({
+ "-1, 1",
+ "31, 1",
+ "3, -1",
+ "3, 31"
+ })
+ void shouldRejectRetryCountsOutsideTheSupportedRange(final int maxRetries,
final int streamingMaxRetries) {
+ assertThrows(IllegalArgumentException.class,
+ () -> new AiProxyRetryProperties(maxRetries,
streamingMaxRetries, Duration.ZERO));
+ }
+
+ @Test
+ void shouldRejectInvalidBackoff() {
+ assertThrows(IllegalArgumentException.class, () -> new
AiProxyRetryProperties(3, 1, null));
+ assertThrows(IllegalArgumentException.class,
+ () -> new AiProxyRetryProperties(3, 1, Duration.ofMillis(-1)));
+ }
+
+ @ParameterizedTest
+ @CsvSource({"-1", "30"})
+ void shouldRejectRetryIndexOutsideTheSupportedRange(final int retryIndex) {
+ final AiProxyRetryProperties properties = new AiProxyRetryProperties();
+
+ assertThrows(IllegalArgumentException.class, () ->
properties.retryDelay(false, retryIndex));
+ }
+}
diff --git
a/shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/shenyu-spring-boot-starter-plugin-ai-proxy/src/main/java/org/apache/shenyu/springboot/starter/plugin/ai/proxy/AiProxyPluginConfiguration.java
b/shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/shenyu-spring-boot-starter-plugin-ai-proxy/src/main/java/org/apache/shenyu/springboot/starter/plugin/ai/proxy/AiProxyPluginConfiguration.java
index a80f1dd00c..d84e90d965 100644
---
a/shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/shenyu-spring-boot-starter-plugin-ai-proxy/src/main/java/org/apache/shenyu/springboot/starter/plugin/ai/proxy/AiProxyPluginConfiguration.java
+++
b/shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/shenyu-spring-boot-starter-plugin-ai-proxy/src/main/java/org/apache/shenyu/springboot/starter/plugin/ai/proxy/AiProxyPluginConfiguration.java
@@ -30,6 +30,7 @@ import
org.apache.shenyu.plugin.ai.proxy.enhanced.provider.OpenAi;
import org.apache.shenyu.plugin.ai.proxy.enhanced.provider.OpenAiCompatible;
import org.apache.shenyu.plugin.ai.proxy.enhanced.service.AiProxyConfigService;
import org.apache.shenyu.plugin.ai.proxy.enhanced.service.AiProxyEngine;
+import
org.apache.shenyu.plugin.ai.proxy.enhanced.service.AiProxyRetryProperties;
import
org.apache.shenyu.plugin.ai.proxy.enhanced.subscriber.CommonAiProxyApiKeyDataSubscriber;
import org.apache.shenyu.plugin.ai.proxy.enhanced.transport.AiProxyTransport;
import
org.apache.shenyu.plugin.ai.proxy.enhanced.transport.AiProxyTransportProperties;
@@ -107,12 +108,15 @@ public class AiProxyPluginConfiguration {
* @param aiProxyProtocolFactory protocol factory
* @param aiProxyProviderFactory provider factory
* @param aiProxyTransport transport
+ * @param retryProperties retry properties
* @return the AI proxy engine
*/
@Bean
public AiProxyEngine aiProxyEngine(final AiProxyProtocolFactory
aiProxyProtocolFactory,
- final AiProxyProviderFactory
aiProxyProviderFactory, final AiProxyTransport aiProxyTransport) {
- return new AiProxyEngine(aiProxyProtocolFactory,
aiProxyProviderFactory, aiProxyTransport);
+ final AiProxyProviderFactory
aiProxyProviderFactory,
+ final AiProxyTransport aiProxyTransport,
+ final AiProxyRetryProperties
retryProperties) {
+ return new AiProxyEngine(aiProxyProtocolFactory,
aiProxyProviderFactory, aiProxyTransport, retryProperties);
}
/**
@@ -244,6 +248,17 @@ public class AiProxyPluginConfiguration {
return new AiProxyTransportProperties();
}
+ /**
+ * AI proxy retry properties.
+ *
+ * @return AI proxy retry properties
+ */
+ @Bean
+ @ConfigurationProperties(prefix = "shenyu.plugins.ai.proxy.retry")
+ public AiProxyRetryProperties aiProxyRetryProperties() {
+ return new AiProxyRetryProperties();
+ }
+
HttpClient createHttpClient(final AiProxyTransportProperties properties) {
return HttpClient.create()
.option(
diff --git
a/shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/shenyu-spring-boot-starter-plugin-ai-proxy/src/test/java/org/apache/shenyu/springboot/starter/plugin/ai/proxy/AiProxyPluginConfigurationTest.java
b/shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/shenyu-spring-boot-starter-plugin-ai-proxy/src/test/java/org/apache/shenyu/springboot/starter/plugin/ai/proxy/AiProxyPluginConfigurationTest.java
index db54707ecf..271edb3cf6 100644
---
a/shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/shenyu-spring-boot-starter-plugin-ai-proxy/src/test/java/org/apache/shenyu/springboot/starter/plugin/ai/proxy/AiProxyPluginConfigurationTest.java
+++
b/shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/shenyu-spring-boot-starter-plugin-ai-proxy/src/test/java/org/apache/shenyu/springboot/starter/plugin/ai/proxy/AiProxyPluginConfigurationTest.java
@@ -32,6 +32,7 @@ import
org.apache.shenyu.plugin.ai.proxy.enhanced.protocol.AiProxyProtocolFactor
import org.apache.shenyu.plugin.ai.proxy.enhanced.protocol.OpenAiResponses;
import
org.apache.shenyu.plugin.ai.proxy.enhanced.provider.AiProxyProviderFactory;
import org.apache.shenyu.plugin.ai.proxy.enhanced.service.AiProxyEngine;
+import
org.apache.shenyu.plugin.ai.proxy.enhanced.service.AiProxyRetryProperties;
import org.apache.shenyu.plugin.ai.proxy.enhanced.transport.AiProxyTransport;
import
org.apache.shenyu.plugin.ai.proxy.enhanced.transport.AiProxyTransportProperties;
import org.junit.jupiter.api.Test;
@@ -128,6 +129,45 @@ class AiProxyPluginConfigurationTest {
});
}
+ @Test
+ void shouldUseDefaultRetryProperties() {
+ contextRunner.run(context -> {
+ final AiProxyRetryProperties properties =
context.getBean(AiProxyRetryProperties.class);
+ assertEquals(3, properties.getMaxRetries());
+ assertEquals(1, properties.getStreamingMaxRetries());
+ assertEquals(Duration.ofSeconds(1),
properties.getInitialBackoff());
+ });
+ }
+
+ @Test
+ void shouldBindConfiguredRetryProperties() {
+ contextRunner
+ .withPropertyValues(
+ "shenyu.plugins.ai.proxy.retry.max-retries=2",
+
"shenyu.plugins.ai.proxy.retry.streaming-max-retries=0",
+ "shenyu.plugins.ai.proxy.retry.initial-backoff=500ms")
+ .run(context -> {
+ final AiProxyRetryProperties properties =
context.getBean(AiProxyRetryProperties.class);
+ assertEquals(2, properties.getMaxRetries());
+ assertEquals(0, properties.getStreamingMaxRetries());
+ assertEquals(Duration.ofMillis(500),
properties.getInitialBackoff());
+ });
+ }
+
+ @Test
+ void shouldRejectInvalidRetryProperties() {
+ for (String property : List.of(
+ "max-retries=-1",
+ "max-retries=31",
+ "streaming-max-retries=-1",
+ "streaming-max-retries=31",
+ "initial-backoff=-1s")) {
+ contextRunner
+ .withPropertyValues("shenyu.plugins.ai.proxy.retry." +
property)
+ .run(context ->
org.assertj.core.api.Assertions.assertThat(context).hasFailed());
+ }
+ }
+
@Test
void shouldConvertConnectTimeoutToNettyMilliseconds() {
final AiProxyTransportProperties properties = new
AiProxyTransportProperties();
