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 34f5180c18 feat(ai-proxy): harden reactive transport (#7629)
34f5180c18 is described below
commit 34f5180c18ba5f8bc763f61ba49601916cf59748
Author: hengyuss <[email protected]>
AuthorDate: Thu Oct 15 10:49:48 2026 +0800
feat(ai-proxy): harden reactive transport (#7629)
---
.../src/main/resources/application.yml | 18 ++
.../plugin/ai/proxy/enhanced/AiProxyPlugin.java | 3 +
.../ai/proxy/enhanced/service/AiProxyEngine.java | 4 +-
.../transport/AiProxyTransportProperties.java | 157 ++++++++++++
.../enhanced/transport/AiTransportException.java | 68 +++++
.../transport/WebClientAiProxyTransport.java | 191 ++++++++++++--
.../ai/proxy/enhanced/AiProxyPluginTest.java | 12 +-
.../transport/WebClientAiProxyTransportTest.java | 285 +++++++++++++++++++++
.../ai/proxy/AiProxyPluginConfiguration.java | 44 +++-
.../ai/proxy/AiProxyPluginConfigurationTest.java | 73 +++++-
10 files changed, 821 insertions(+), 34 deletions(-)
diff --git a/shenyu-bootstrap/src/main/resources/application.yml
b/shenyu-bootstrap/src/main/resources/application.yml
index e861868a07..cdd1e4b9a3 100644
--- a/shenyu-bootstrap/src/main/resources/application.yml
+++ b/shenyu-bootstrap/src/main/resources/application.yml
@@ -308,6 +308,24 @@ shenyu:
props:
jvm_enabled: true
# plugins:
+# ai:
+# proxy:
+# transport:
+# # Maximum time allowed to establish a connection to the AI upstream.
+# # Default: 30s.
+# connect-timeout: 30s
+# # Maximum time to wait for upstream HTTP response headers.
+# # It does not limit the subsequent SSE body. Default: 30s.
+# response-header-timeout: 30s
+# # Maximum interval between two upstream SSE or body chunks.
+# # The timer restarts after each chunk. Default: 60s.
+# streaming-idle-timeout: 60s
+# # Maximum duration of the entire upstream response body stream.
+# # New chunks do not reset this limit. Default: 10m.
+# total-stream-duration: 10m
+# # Maximum cumulative bytes read from one upstream response.
+# # Zero means unlimited. Default: 0.
+# max-response-bytes: 0
# rate-limiter.enabled: false
local:
enabled: false
diff --git
a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/main/java/org/apache/shenyu/plugin/ai/proxy/enhanced/AiProxyPlugin.java
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/main/java/org/apache/shenyu/plugin/ai/proxy/enhanced/AiProxyPlugin.java
index f3eb14c346..31bcb6d2d4 100644
---
a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/main/java/org/apache/shenyu/plugin/ai/proxy/enhanced/AiProxyPlugin.java
+++
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/main/java/org/apache/shenyu/plugin/ai/proxy/enhanced/AiProxyPlugin.java
@@ -30,6 +30,7 @@ import
org.apache.shenyu.plugin.ai.proxy.enhanced.model.AiProxyException;
import org.apache.shenyu.plugin.ai.proxy.enhanced.model.AiUpstreamResponse;
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.transport.AiTransportException;
import org.apache.shenyu.plugin.api.ShenyuPluginChain;
import org.apache.shenyu.plugin.base.AbstractShenyuPlugin;
import org.apache.shenyu.plugin.base.utils.CacheKeyUtils;
@@ -139,6 +140,8 @@ public class AiProxyPlugin extends AbstractShenyuPlugin {
status = HttpStatus.PAYLOAD_TOO_LARGE;
} else if (error instanceof ShenyuException) {
status = HttpStatus.BAD_REQUEST;
+ } else if (error instanceof AiTransportException) {
+ status = HttpStatus.BAD_GATEWAY;
} else if (error instanceof AiProxyException) {
status = ((AiProxyException) error).getKind() ==
AiProxyException.Kind.REQUEST ? HttpStatus.BAD_REQUEST : HttpStatus.BAD_GATEWAY;
}
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 b25664b097..42a0811487 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
@@ -33,6 +33,7 @@ import
org.apache.shenyu.plugin.ai.proxy.enhanced.protocol.AiProxyProtocolFactor
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.transport.AiProxyTransport;
+import
org.apache.shenyu.plugin.ai.proxy.enhanced.transport.AiTransportException;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
@@ -128,7 +129,8 @@ public final class AiProxyEngine {
}).then();
})
.onErrorResume(error -> {
- boolean transportFailure = error instanceof
AiProxyException
+ boolean transportFailure = error instanceof
AiTransportException
+ || error instanceof AiProxyException
&& ((AiProxyException) error).getKind() ==
AiProxyException.Kind.TRANSPORT;
boolean statusFailure = error instanceof
UpstreamStatusFailure;
if (emitted.get() || !(transportFailure || statusFailure))
{
diff --git
a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/main/java/org/apache/shenyu/plugin/ai/proxy/enhanced/transport/AiProxyTransportProperties.java
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/main/java/org/apache/shenyu/plugin/ai/proxy/enhanced/transport/AiProxyTransportProperties.java
new file mode 100644
index 0000000000..c76939b0f1
--- /dev/null
+++
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/main/java/org/apache/shenyu/plugin/ai/proxy/enhanced/transport/AiProxyTransportProperties.java
@@ -0,0 +1,157 @@
+/*
+ * 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.transport;
+
+import java.time.Duration;
+import java.util.Objects;
+
+/**
+ * AI proxy transport properties.
+ */
+public class AiProxyTransportProperties {
+
+ /**
+ * Maximum duration allowed for establishing an upstream connection.
+ */
+ private Duration connectTimeout = Duration.ofSeconds(30);
+
+ /**
+ * Maximum duration allowed for receiving upstream response headers.
+ */
+ private Duration responseHeaderTimeout = Duration.ofSeconds(30);
+
+ /**
+ * Maximum idle duration allowed between upstream response body chunks.
+ */
+ private Duration streamingIdleTimeout = Duration.ofSeconds(60);
+
+ /**
+ * Maximum total duration allowed for an upstream response body stream.
+ */
+ private Duration totalStreamDuration = Duration.ofMinutes(10);
+
+ /**
+ * Maximum total bytes allowed for one upstream response. Zero means
unlimited.
+ */
+ private long maxResponseBytes;
+
+ /**
+ * Get connect timeout.
+ *
+ * @return connect timeout
+ */
+ public Duration getConnectTimeout() {
+ return connectTimeout;
+ }
+
+ /**
+ * Set connect timeout.
+ *
+ * @param connectTimeout connect timeout
+ */
+ public void setConnectTimeout(final Duration connectTimeout) {
+ this.connectTimeout = requirePositive(connectTimeout,
"connectTimeout");
+ }
+
+ /**
+ * Get response header timeout.
+ *
+ * @return response header timeout
+ */
+ public Duration getResponseHeaderTimeout() {
+ return responseHeaderTimeout;
+ }
+
+ /**
+ * Set response header timeout.
+ *
+ * @param responseHeaderTimeout response header timeout
+ */
+ public void setResponseHeaderTimeout(final Duration responseHeaderTimeout)
{
+ this.responseHeaderTimeout = requirePositive(responseHeaderTimeout,
"responseHeaderTimeout");
+ }
+
+ /**
+ * Get streaming idle timeout.
+ *
+ * @return streaming idle timeout
+ */
+ public Duration getStreamingIdleTimeout() {
+ return streamingIdleTimeout;
+ }
+
+ /**
+ * Set streaming idle timeout.
+ *
+ * @param streamingIdleTimeout streaming idle timeout
+ */
+ public void setStreamingIdleTimeout(final Duration streamingIdleTimeout) {
+ this.streamingIdleTimeout = requirePositive(streamingIdleTimeout,
"streamingIdleTimeout");
+ }
+
+ /**
+ * Get total stream duration.
+ *
+ * @return total stream duration
+ */
+ public Duration getTotalStreamDuration() {
+ return totalStreamDuration;
+ }
+
+ /**
+ * Set total stream duration.
+ *
+ * @param totalStreamDuration total stream duration
+ */
+ public void setTotalStreamDuration(final Duration totalStreamDuration) {
+ this.totalStreamDuration = requirePositive(totalStreamDuration,
"totalStreamDuration");
+ }
+
+ /**
+ * Get maximum upstream response bytes.
+ *
+ * @return maximum upstream response bytes, or zero when unlimited
+ */
+ public long getMaxResponseBytes() {
+ return maxResponseBytes;
+ }
+
+ /**
+ * Set maximum upstream response bytes.
+ *
+ * @param maxResponseBytes maximum upstream response bytes, or zero when
unlimited
+ */
+ public void setMaxResponseBytes(final long maxResponseBytes) {
+ this.maxResponseBytes = requireNonNegative(maxResponseBytes,
"maxResponseBytes");
+ }
+
+ static Duration requirePositive(final Duration duration, final String
name) {
+ Objects.requireNonNull(duration, name + " must not be null");
+ if (duration.isZero() || duration.isNegative()) {
+ throw new IllegalArgumentException(name + " must be greater than
zero");
+ }
+ return duration;
+ }
+
+ static long requireNonNegative(final long value, final String name) {
+ if (value < 0) {
+ throw new IllegalArgumentException(name + " must be greater than
or equal to zero");
+ }
+ return value;
+ }
+}
diff --git
a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/main/java/org/apache/shenyu/plugin/ai/proxy/enhanced/transport/AiTransportException.java
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/main/java/org/apache/shenyu/plugin/ai/proxy/enhanced/transport/AiTransportException.java
new file mode 100644
index 0000000000..aed6e48b33
--- /dev/null
+++
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/main/java/org/apache/shenyu/plugin/ai/proxy/enhanced/transport/AiTransportException.java
@@ -0,0 +1,68 @@
+/*
+ * 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.transport;
+
+/**
+ * Classified failure raised by an AI transport.
+ */
+public final class AiTransportException extends RuntimeException {
+
+ private static final long serialVersionUID = 1L;
+
+ private final Kind kind;
+
+ /**
+ * Create an AI transport exception.
+ *
+ * @param kind transport failure kind
+ * @param message failure message
+ * @param cause original failure
+ */
+ public AiTransportException(
+ final Kind kind,
+ final String message,
+ final Throwable cause) {
+ super(message, cause);
+ this.kind = kind;
+ }
+
+ /**
+ * Get transport failure kind.
+ *
+ * @return transport failure kind
+ */
+ public Kind getKind() {
+ return kind;
+ }
+
+ /**
+ * Transport failure classification.
+ */
+ public enum Kind {
+ CONNECT_TIMEOUT,
+ RESPONSE_HEADER_TIMEOUT,
+ STREAM_IDLE_TIMEOUT,
+ STREAM_DURATION_TIMEOUT,
+ RESPONSE_TOO_LARGE,
+ DNS_FAILURE,
+ TLS_FAILURE,
+ CONNECTION_RESET,
+ IO_FAILURE,
+ UNKNOWN
+ }
+}
diff --git
a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/main/java/org/apache/shenyu/plugin/ai/proxy/enhanced/transport/WebClientAiProxyTransport.java
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/main/java/org/apache/shenyu/plugin/ai/proxy/enhanced/transport/WebClientAiProxyTransport.java
index 090966023f..06be8f997f 100644
---
a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/main/java/org/apache/shenyu/plugin/ai/proxy/enhanced/transport/WebClientAiProxyTransport.java
+++
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/main/java/org/apache/shenyu/plugin/ai/proxy/enhanced/transport/WebClientAiProxyTransport.java
@@ -17,16 +17,26 @@
package org.apache.shenyu.plugin.ai.proxy.enhanced.transport;
+import io.netty.channel.ConnectTimeoutException;
import org.apache.shenyu.plugin.ai.proxy.enhanced.model.AiUpstreamRequest;
import org.apache.shenyu.plugin.ai.proxy.enhanced.model.AiUpstreamResponse;
-import org.springframework.http.HttpMethod;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
-import org.apache.shenyu.plugin.ai.proxy.enhanced.model.AiProxyException;
+import org.springframework.http.HttpHeaders;
+import org.springframework.http.HttpMethod;
import org.springframework.web.reactive.function.client.WebClient;
+import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
+import reactor.netty.http.client.PrematureCloseException;
+import javax.net.ssl.SSLException;
+import java.io.IOException;
+import java.net.UnknownHostException;
+import java.time.Duration;
import java.util.Objects;
+import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicLong;
/**
* AI transport implemented with Spring WebClient.
@@ -35,13 +45,40 @@ public final class WebClientAiProxyTransport implements
AiProxyTransport {
private final WebClient webClient;
+ private final Duration responseHeaderTimeout;
+
+ private final Duration streamingIdleTimeout;
+
+ private final Duration totalStreamDuration;
+
+ private final long maxResponseBytes;
+
/**
* Create a WebClient transport.
*
* @param webClient web client
*/
public WebClientAiProxyTransport(final WebClient webClient) {
- this.webClient = webClient;
+ this(webClient, new AiProxyTransportProperties());
+ }
+
+ /**
+ * Create a WebClient transport.
+ *
+ * @param webClient web client
+ * @param properties transport properties
+ */
+ public WebClientAiProxyTransport(final WebClient webClient, final
AiProxyTransportProperties properties) {
+ this.webClient = Objects.requireNonNull(webClient, "webClient must not
be null");
+ Objects.requireNonNull(properties, "properties must not be null");
+ this.responseHeaderTimeout =
AiProxyTransportProperties.requirePositive(
+ properties.getResponseHeaderTimeout(),
"responseHeaderTimeout");
+ this.streamingIdleTimeout = AiProxyTransportProperties.requirePositive(
+ properties.getStreamingIdleTimeout(), "streamingIdleTimeout");
+ this.totalStreamDuration = AiProxyTransportProperties.requirePositive(
+ properties.getTotalStreamDuration(), "totalStreamDuration");
+ this.maxResponseBytes = AiProxyTransportProperties.requireNonNegative(
+ properties.getMaxResponseBytes(), "maxResponseBytes");
}
@Override
@@ -65,20 +102,138 @@ public final class WebClientAiProxyTransport implements
AiProxyTransport {
.retrieve()
.onRawStatus(status -> status >= 400, response -> Mono.empty())
.toEntityFlux(DataBuffer.class)
- .map(entity -> new AiUpstreamResponse(
- entity.getStatusCode().value(),
- entity.getHeaders(),
- Objects.requireNonNull(entity.getBody()).map(buffer ->
{
- try {
- byte[] bytes = new
byte[buffer.readableByteCount()];
- buffer.read(bytes);
- return bytes;
- } finally {
- DataBufferUtils.release(buffer);
- }
- }).doOnDiscard(DataBuffer.class,
DataBufferUtils::release)
- .onErrorMap(error -> new
AiProxyException(AiProxyException.Kind.TRANSPORT, "AI upstream body failed",
error))))
- .onErrorMap(error -> !(error instanceof AiProxyException),
- error -> new
AiProxyException(AiProxyException.Kind.TRANSPORT, "AI upstream exchange
failed", error));
+ .timeout(responseHeaderTimeout)
+ .onErrorMap(TimeoutException.class, error -> new
AiTransportException(
+ AiTransportException.Kind.RESPONSE_HEADER_TIMEOUT,
+ "AI upstream response header timed out",
+ error))
+ .map(entity -> {
+ final Flux<DataBuffer> rawBody =
Objects.requireNonNull(entity.getBody());
+ if (exceedsMaxResponseBytes(entity.getHeaders())) {
+ return new AiUpstreamResponse(
+ entity.getStatusCode().value(),
+ entity.getHeaders(),
+ rejectOversizedResponse(rawBody));
+ }
+
+ final Flux<DataBuffer> dataBuffers =
enforceTotalStreamDuration(
+ rawBody.timeout(streamingIdleTimeout)
+ .onErrorMap(TimeoutException.class, error
-> new AiTransportException(
+
AiTransportException.Kind.STREAM_IDLE_TIMEOUT,
+ "AI upstream response body became
idle",
+ error)));
+ final Flux<byte[]> bytes = dataBuffers
+ .map(buffer -> {
+ try {
+ final byte[] result = new
byte[buffer.readableByteCount()];
+ buffer.read(result);
+ return result;
+ } finally {
+ DataBufferUtils.release(buffer);
+ }
+ })
+ .doOnDiscard(DataBuffer.class,
DataBufferUtils::release)
+ .onErrorMap(error -> !(error instanceof
AiTransportException),
+ error -> toTransportException("AI upstream
body failed", error));
+ return new AiUpstreamResponse(
+ entity.getStatusCode().value(),
+ entity.getHeaders(),
+ enforceMaxResponseBytes(bytes));
+ })
+ .onErrorMap(error -> !(error instanceof AiTransportException),
+ error -> toTransportException("AI upstream exchange
failed", error));
+ }
+
+ private Flux<DataBuffer> enforceTotalStreamDuration(final Flux<DataBuffer>
body) {
+ return Flux.defer(() -> {
+ final AtomicBoolean expired = new AtomicBoolean();
+ final Mono<Long> deadline = Mono.delay(totalStreamDuration)
+ .doOnNext(ignored -> expired.set(true));
+ return body.takeUntilOther(deadline)
+ .concatWith(Flux.defer(() -> expired.get()
+ ? Flux.error(new AiTransportException(
+
AiTransportException.Kind.STREAM_DURATION_TIMEOUT,
+ "AI upstream response body exceeded total
duration",
+ new TimeoutException("AI upstream response
body exceeded total duration")))
+ : Flux.empty()));
+ });
+ }
+
+ private AiTransportException toTransportException(final String message,
final Throwable failure) {
+ return new AiTransportException(classify(failure), message, failure);
+ }
+
+ private AiTransportException.Kind classify(final Throwable failure) {
+ if (hasCause(failure, ConnectTimeoutException.class)) {
+ return AiTransportException.Kind.CONNECT_TIMEOUT;
+ }
+ if (hasCause(failure, UnknownHostException.class)) {
+ return AiTransportException.Kind.DNS_FAILURE;
+ }
+ if (hasCause(failure, SSLException.class)) {
+ return AiTransportException.Kind.TLS_FAILURE;
+ }
+ if (hasCause(failure, PrematureCloseException.class)) {
+ return AiTransportException.Kind.CONNECTION_RESET;
+ }
+ if (hasCause(failure, IOException.class)) {
+ return AiTransportException.Kind.IO_FAILURE;
+ }
+ return AiTransportException.Kind.UNKNOWN;
+ }
+
+ private boolean hasCause(final Throwable failure, final Class<? extends
Throwable> type) {
+ Throwable current = failure;
+ while (Objects.nonNull(current)) {
+ if (type.isInstance(current)) {
+ return true;
+ }
+ Throwable cause = current.getCause();
+ if (cause == current) {
+ return false;
+ }
+ current = cause;
+ }
+ return false;
+ }
+
+ private Flux<byte[]> enforceMaxResponseBytes(final Flux<byte[]> body) {
+ if (maxResponseBytes == 0) {
+ return body;
+ }
+ return Flux.defer(() -> {
+ final AtomicLong receivedBytes = new AtomicLong();
+ return body.handle((bytes, sink) -> {
+ final long received = receivedBytes.get();
+ if (bytes.length > maxResponseBytes - received) {
+ sink.error(responseTooLarge());
+ return;
+ }
+ receivedBytes.set(received + bytes.length);
+ sink.next(bytes);
+ });
+ });
+ }
+
+ private boolean exceedsMaxResponseBytes(
+ final HttpHeaders headers) {
+ if (maxResponseBytes == 0) {
+ return false;
+ }
+
+ final long contentLength = headers.getContentLength();
+ return contentLength >= 0
+ && contentLength > maxResponseBytes;
+ }
+
+ private Flux<byte[]> rejectOversizedResponse(final Flux<DataBuffer> body) {
+ return body.take(0, false).thenMany(Flux.error(responseTooLarge()));
+ }
+
+ private AiTransportException responseTooLarge() {
+ return new AiTransportException(
+ AiTransportException.Kind.RESPONSE_TOO_LARGE,
+ "AI upstream response exceeded maximum size",
+ null);
}
}
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 2a46f0b46a..bda2653b1f 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
@@ -27,7 +27,6 @@ import
org.apache.shenyu.common.dto.convert.rule.AiProxyHandle;
import org.apache.shenyu.common.enums.PluginEnum;
import org.apache.shenyu.plugin.ai.proxy.enhanced.cache.AiProxyApiKeyCache;
import org.apache.shenyu.plugin.ai.proxy.enhanced.handler.AiProxyPluginHandler;
-import org.apache.shenyu.plugin.ai.proxy.enhanced.model.AiProxyException;
import org.apache.shenyu.plugin.ai.proxy.enhanced.model.AiUpstreamRequest;
import org.apache.shenyu.plugin.ai.proxy.enhanced.model.AiUpstreamResponse;
import
org.apache.shenyu.plugin.ai.proxy.enhanced.protocol.AiProxyProtocolFactory;
@@ -38,7 +37,7 @@ 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.transport.AiProxyTransport;
-import org.apache.shenyu.plugin.api.ShenyuPluginChain;
+import
org.apache.shenyu.plugin.ai.proxy.enhanced.transport.AiTransportException;
import org.apache.shenyu.plugin.base.utils.CacheKeyUtils;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
@@ -63,7 +62,6 @@ import java.util.concurrent.atomic.AtomicBoolean;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertTrue;
-import static org.mockito.Mockito.mock;
/**
* Compatibility tests through AiProxyEngine instead of Spring AI clients.
@@ -247,7 +245,7 @@ class AiProxyPluginTest {
}).then());
AiProxyPlugin plugin = plugin(request -> Mono.just(new
AiUpstreamResponse(200, Map.of(),
Flux.just(bytes("data:
{\"first\":true}\n\n")).concatWith(Flux.error(failure())))));
-
StepVerifier.create(execute(plugin)).expectError(AiProxyException.class).verify();
+
StepVerifier.create(execute(plugin)).expectError(AiTransportException.class).verify();
assertEquals(1, requests.size());
assertEquals(List.of("data: {\"first\":true}\n\n"), received);
assertEquals(HttpStatus.OK, exchange.getResponse().getStatusCode());
@@ -289,7 +287,7 @@ class AiProxyPluginTest {
}
private Mono<Void> execute(final AiProxyPlugin plugin) {
- return plugin.doExecute(exchange, mock(ShenyuPluginChain.class),
selector, rule);
+ return plugin.doExecute(exchange, current -> Mono.empty(), selector,
rule);
}
private AiUpstreamResponse response(final int status, final String body) {
@@ -309,7 +307,7 @@ class AiProxyPluginTest {
return config;
}
- private AiProxyException failure() {
- return new AiProxyException(AiProxyException.Kind.TRANSPORT, "fixture
failure", null);
+ private AiTransportException failure() {
+ return new AiTransportException(AiTransportException.Kind.UNKNOWN,
"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/transport/WebClientAiProxyTransportTest.java
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/test/java/org/apache/shenyu/plugin/ai/proxy/enhanced/transport/WebClientAiProxyTransportTest.java
index 2583ac98ca..9625024719 100644
---
a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/test/java/org/apache/shenyu/plugin/ai/proxy/enhanced/transport/WebClientAiProxyTransportTest.java
+++
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy/src/test/java/org/apache/shenyu/plugin/ai/proxy/enhanced/transport/WebClientAiProxyTransportTest.java
@@ -17,9 +17,11 @@
package org.apache.shenyu.plugin.ai.proxy.enhanced.transport;
+import io.netty.channel.ConnectTimeoutException;
import org.apache.shenyu.plugin.ai.proxy.enhanced.model.AiUpstreamRequest;
import org.apache.shenyu.plugin.ai.proxy.enhanced.model.AiUpstreamResponse;
import org.junit.jupiter.api.Test;
+import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DefaultDataBufferFactory;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpStatus;
@@ -29,15 +31,24 @@ import
org.springframework.web.reactive.function.client.ClientResponse;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
+import reactor.netty.http.client.PrematureCloseException;
import reactor.test.StepVerifier;
+import javax.net.ssl.SSLException;
+import java.io.IOException;
import java.net.URI;
+import java.net.UnknownHostException;
import java.nio.charset.StandardCharsets;
+import java.time.Duration;
import java.util.List;
import java.util.Map;
+import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
class WebClientAiProxyTransportTest {
@@ -87,6 +98,258 @@ class WebClientAiProxyTransportTest {
.verifyComplete();
}
+ @Test
+ void testAllowsResponseExactlyAtMaximumSize() {
+ final WebClient webClient =
webClientWithBody(Flux.just(dataBuffer("ab"), dataBuffer("cde")));
+ final AiProxyTransportProperties properties = new
AiProxyTransportProperties();
+ properties.setMaxResponseBytes(5);
+ final WebClientAiProxyTransport transport = new
WebClientAiProxyTransport(webClient, properties);
+
+
StepVerifier.create(transport.execute(createRequest()).flatMapMany(AiUpstreamResponse::getBody)
+ .map(bytes -> new String(bytes,
StandardCharsets.UTF_8)))
+ .expectNext("ab", "cde")
+ .verifyComplete();
+ }
+
+ @Test
+ void testAllowsContentLengthExactlyAtMaximumSize() {
+ final WebClient webClient = WebClient.builder()
+ .exchangeFunction(request ->
Mono.just(ClientResponse.create(HttpStatus.OK)
+ .header(HttpHeaders.CONTENT_LENGTH, "5")
+ .body(Flux.just(dataBuffer("abcde")))
+ .build()))
+ .build();
+ final AiProxyTransportProperties properties = new
AiProxyTransportProperties();
+ properties.setMaxResponseBytes(5);
+ final WebClientAiProxyTransport transport = new
WebClientAiProxyTransport(webClient, properties);
+
+
StepVerifier.create(transport.execute(createRequest()).flatMapMany(AiUpstreamResponse::getBody)
+ .map(bytes -> new String(bytes,
StandardCharsets.UTF_8)))
+ .expectNext("abcde")
+ .verifyComplete();
+ }
+
+ @Test
+ void testRejectsOversizedContentLengthAndCancelsBody() {
+ final AtomicBoolean subscribed = new AtomicBoolean();
+ final AtomicBoolean cancelled = new AtomicBoolean();
+ final Flux<DataBuffer> body = Flux.<DataBuffer>never()
+ .doOnSubscribe(ignored -> subscribed.set(true))
+ .doOnCancel(() -> cancelled.set(true));
+ final WebClient webClient = WebClient.builder()
+ .exchangeFunction(request ->
Mono.just(ClientResponse.create(HttpStatus.OK)
+ .header(HttpHeaders.CONTENT_LENGTH, "6")
+ .body(body)
+ .build()))
+ .build();
+ final AiProxyTransportProperties properties = new
AiProxyTransportProperties();
+ properties.setMaxResponseBytes(5);
+ final WebClientAiProxyTransport transport = new
WebClientAiProxyTransport(webClient, properties);
+
+
StepVerifier.create(transport.execute(createRequest()).flatMapMany(AiUpstreamResponse::getBody))
+ .expectErrorSatisfies(error ->
assertEquals(AiTransportException.Kind.RESPONSE_TOO_LARGE,
+ assertInstanceOf(AiTransportException.class,
error).getKind()))
+ .verify();
+ assertTrue(subscribed.get());
+ assertTrue(cancelled.get());
+ }
+
+ @Test
+ void testRejectsChunkThatExceedsCumulativeResponseSize() {
+ final AtomicBoolean cancelled = new AtomicBoolean();
+ final Flux<DataBuffer> body = Flux.concat(
+ Flux.just(dataBuffer("abc"), dataBuffer("def")),
+ Flux.<DataBuffer>never())
+ .doOnCancel(() -> cancelled.set(true));
+ final WebClient webClient = webClientWithBody(body);
+ final AiProxyTransportProperties properties = new
AiProxyTransportProperties();
+ properties.setMaxResponseBytes(5);
+ final WebClientAiProxyTransport transport = new
WebClientAiProxyTransport(webClient, properties);
+
+
StepVerifier.create(transport.execute(createRequest()).flatMapMany(AiUpstreamResponse::getBody)
+ .map(bytes -> new String(bytes,
StandardCharsets.UTF_8)))
+ .expectNext("abc")
+ .expectErrorSatisfies(error ->
assertEquals(AiTransportException.Kind.RESPONSE_TOO_LARGE,
+ assertInstanceOf(AiTransportException.class,
error).getKind()))
+ .verify();
+ assertTrue(cancelled.get());
+ }
+
+ @Test
+ void testRejectsOversizedFirstChunkWithoutEmittingIt() {
+ final WebClient webClient =
webClientWithBody(Flux.just(dataBuffer("abcdef")));
+ final AiProxyTransportProperties properties = new
AiProxyTransportProperties();
+ properties.setMaxResponseBytes(5);
+ final WebClientAiProxyTransport transport = new
WebClientAiProxyTransport(webClient, properties);
+
+
StepVerifier.create(transport.execute(createRequest()).flatMapMany(AiUpstreamResponse::getBody))
+ .expectErrorSatisfies(error ->
assertEquals(AiTransportException.Kind.RESPONSE_TOO_LARGE,
+ assertInstanceOf(AiTransportException.class,
error).getKind()))
+ .verify();
+ }
+
+ @Test
+ void testClassifiesExchangeFailures() {
+ assertFailureKind(new ConnectTimeoutException(),
AiTransportException.Kind.CONNECT_TIMEOUT);
+ assertFailureKind(new UnknownHostException(),
AiTransportException.Kind.DNS_FAILURE);
+ assertFailureKind(new SSLException("handshake failed"),
AiTransportException.Kind.TLS_FAILURE);
+ assertFailureKind(PrematureCloseException.TEST_EXCEPTION,
AiTransportException.Kind.CONNECTION_RESET);
+ assertFailureKind(new IOException("read failed"),
AiTransportException.Kind.IO_FAILURE);
+ assertFailureKind(new IllegalStateException("unexpected"),
AiTransportException.Kind.UNKNOWN);
+ }
+
+ @Test
+ void testClassifiesWrappedBodyFailure() {
+ final WebClient webClient = WebClient.builder()
+ .exchangeFunction(request ->
Mono.just(ClientResponse.create(HttpStatus.OK)
+ .body(Flux.error(new IllegalStateException(new
SSLException("handshake failed"))))
+ .build()))
+ .build();
+ final AiUpstreamResponse response = new
WebClientAiProxyTransport(webClient).execute(createRequest()).block();
+
+ StepVerifier.create(response.getBody())
+ .expectErrorSatisfies(error ->
assertEquals(AiTransportException.Kind.TLS_FAILURE,
+ assertInstanceOf(AiTransportException.class,
error).getKind()))
+ .verify();
+ }
+
+ @Test
+ void testResponseHeaderTimeout() {
+ final WebClient webClient = WebClient.builder()
+ .exchangeFunction(request -> Mono.delay(Duration.ofSeconds(10))
+
.thenReturn(ClientResponse.create(HttpStatus.OK).build()))
+ .build();
+ final WebClientAiProxyTransport transport = new
WebClientAiProxyTransport(
+ webClient, properties(Duration.ofSeconds(5),
Duration.ofSeconds(60), Duration.ofMinutes(10)));
+
+ StepVerifier.withVirtualTime(() -> transport.execute(createRequest()))
+ .thenAwait(Duration.ofSeconds(5))
+ .expectErrorSatisfies(error ->
assertEquals(AiTransportException.Kind.RESPONSE_HEADER_TIMEOUT,
+ assertInstanceOf(AiTransportException.class,
error).getKind()))
+ .verify();
+ }
+
+ @Test
+ void testResponseHeaderTimeoutDoesNotLimitBodyStream() {
+ final WebClient webClient = WebClient.builder()
+ .exchangeFunction(request ->
Mono.just(ClientResponse.create(HttpStatus.OK)
+ .body(Flux.never())
+ .build()))
+ .build();
+ final WebClientAiProxyTransport transport = new
WebClientAiProxyTransport(
+ webClient, properties(Duration.ofSeconds(5),
Duration.ofSeconds(60), Duration.ofMinutes(10)));
+
+ StepVerifier.withVirtualTime(() -> transport.execute(createRequest()))
+ .expectNextMatches(response -> response.getStatusCode() ==
HttpStatus.OK.value())
+ .verifyComplete();
+ }
+
+ @Test
+ void testStreamingIdleTimeout() {
+ final WebClient webClient = WebClient.builder()
+ .exchangeFunction(request ->
Mono.just(ClientResponse.create(HttpStatus.OK)
+ .body(Flux.never())
+ .build()))
+ .build();
+ final WebClientAiProxyTransport transport = new
WebClientAiProxyTransport(
+ webClient, properties(Duration.ofSeconds(5),
Duration.ofSeconds(10), Duration.ofMinutes(10)));
+
+ StepVerifier.withVirtualTime(() -> transport.execute(createRequest())
+ .flatMapMany(AiUpstreamResponse::getBody))
+ .thenAwait(Duration.ofSeconds(10))
+ .expectErrorSatisfies(error ->
assertEquals(AiTransportException.Kind.STREAM_IDLE_TIMEOUT,
+ assertInstanceOf(AiTransportException.class,
error).getKind()))
+ .verify();
+ }
+
+ @Test
+ void testStreamingIdleTimeoutRestartsAfterEachChunk() {
+ final Flux<DataBuffer> body = Flux.defer(() ->
Flux.interval(Duration.ofSeconds(4))
+ .take(2)
+ .map(index -> DefaultDataBufferFactory.sharedInstance.wrap(
+ ("chunk-" + index).getBytes(StandardCharsets.UTF_8))));
+ final WebClient webClient = WebClient.builder()
+ .exchangeFunction(request ->
Mono.just(ClientResponse.create(HttpStatus.OK)
+ .body(body)
+ .build()))
+ .build();
+ final WebClientAiProxyTransport transport = new
WebClientAiProxyTransport(
+ webClient, properties(Duration.ofSeconds(5),
Duration.ofSeconds(5), Duration.ofMinutes(10)));
+
+ StepVerifier.withVirtualTime(() -> transport.execute(createRequest())
+ .flatMapMany(AiUpstreamResponse::getBody)
+ .map(bytes -> new String(bytes,
StandardCharsets.UTF_8)))
+ .thenAwait(Duration.ofSeconds(8))
+ .expectNext("chunk-0", "chunk-1")
+ .verifyComplete();
+ }
+
+ @Test
+ void testTotalStreamDurationCancelsActiveStream() {
+ final AtomicBoolean cancelled = new AtomicBoolean();
+ final Flux<DataBuffer> body = Flux.defer(() ->
Flux.interval(Duration.ofSeconds(1))
+ .<DataBuffer>map(index ->
DefaultDataBufferFactory.sharedInstance.wrap(
+ ("chunk-" + index).getBytes(StandardCharsets.UTF_8))))
+ .doOnCancel(() -> cancelled.set(true));
+ final WebClient webClient = WebClient.builder()
+ .exchangeFunction(request ->
Mono.just(ClientResponse.create(HttpStatus.OK)
+ .body(body)
+ .build()))
+ .build();
+ final WebClientAiProxyTransport transport = new
WebClientAiProxyTransport(
+ webClient, properties(Duration.ofSeconds(5),
Duration.ofSeconds(2), Duration.ofMillis(4500)));
+
+ StepVerifier.withVirtualTime(() -> transport.execute(createRequest())
+ .flatMapMany(AiUpstreamResponse::getBody))
+ .thenAwait(Duration.ofMillis(4500))
+ .expectNextCount(4)
+ .expectErrorSatisfies(error ->
assertEquals(AiTransportException.Kind.STREAM_DURATION_TIMEOUT,
+ assertInstanceOf(AiTransportException.class,
error).getKind()))
+ .verify();
+ assertEquals(true, cancelled.get());
+ }
+
+ @Test
+ void testTotalStreamDurationAllowsNormalCompletion() {
+ final Flux<DataBuffer> body =
Flux.just(DefaultDataBufferFactory.sharedInstance.wrap(
+ "complete".getBytes(StandardCharsets.UTF_8)));
+ final WebClient webClient = WebClient.builder()
+ .exchangeFunction(request ->
Mono.just(ClientResponse.create(HttpStatus.OK)
+ .body(body)
+ .build()))
+ .build();
+ final WebClientAiProxyTransport transport = new
WebClientAiProxyTransport(
+ webClient, properties(Duration.ofSeconds(5),
Duration.ofSeconds(5), Duration.ofSeconds(10)));
+
+ StepVerifier.withVirtualTime(() -> transport.execute(createRequest())
+ .flatMapMany(AiUpstreamResponse::getBody)
+ .map(bytes -> new String(bytes,
StandardCharsets.UTF_8)))
+ .expectNext("complete")
+ .verifyComplete();
+ }
+
+ @Test
+ void testRejectsNonPositiveTimeouts() {
+ final AiProxyTransportProperties properties = new
AiProxyTransportProperties();
+
+ assertThrows(IllegalArgumentException.class, () ->
properties.setResponseHeaderTimeout(Duration.ZERO));
+ assertThrows(IllegalArgumentException.class, () ->
properties.setStreamingIdleTimeout(Duration.ofSeconds(-1)));
+ assertThrows(IllegalArgumentException.class, () ->
properties.setTotalStreamDuration(Duration.ZERO));
+ assertThrows(IllegalArgumentException.class, () ->
properties.setMaxResponseBytes(-1));
+ }
+
+ private void assertFailureKind(final Throwable failure, final
AiTransportException.Kind expected) {
+ final WebClient webClient = WebClient.builder()
+ .exchangeFunction(request -> Mono.error(failure))
+ .build();
+ final WebClientAiProxyTransport transport = new
WebClientAiProxyTransport(webClient);
+
+ StepVerifier.create(transport.execute(createRequest()))
+ .expectErrorSatisfies(error -> assertEquals(expected,
+ assertInstanceOf(AiTransportException.class,
error).getKind()))
+ .verify();
+ }
+
private AiUpstreamRequest createRequest() {
final AiUpstreamRequest request = new AiUpstreamRequest();
request.setMethod("POST");
@@ -97,4 +360,26 @@ class WebClientAiProxyTransportTest {
request.setBody("{}".getBytes(StandardCharsets.UTF_8));
return request;
}
+
+ private WebClient webClientWithBody(final Flux<DataBuffer> body) {
+ return WebClient.builder()
+ .exchangeFunction(request ->
Mono.just(ClientResponse.create(HttpStatus.OK)
+ .body(body)
+ .build()))
+ .build();
+ }
+
+ private DataBuffer dataBuffer(final String value) {
+ return
DefaultDataBufferFactory.sharedInstance.wrap(value.getBytes(StandardCharsets.UTF_8));
+ }
+
+ private AiProxyTransportProperties properties(final Duration
responseHeaderTimeout,
+ final Duration
streamingIdleTimeout,
+ final Duration
totalStreamDuration) {
+ final AiProxyTransportProperties properties = new
AiProxyTransportProperties();
+ properties.setResponseHeaderTimeout(responseHeaderTimeout);
+ properties.setStreamingIdleTimeout(streamingIdleTimeout);
+ properties.setTotalStreamDuration(totalStreamDuration);
+ return properties;
+ }
}
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 1708b69a2e..07bf852251 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
@@ -17,6 +17,7 @@
package org.apache.shenyu.springboot.starter.plugin.ai.proxy;
+import io.netty.channel.ChannelOption;
import org.apache.shenyu.plugin.ai.proxy.enhanced.AiProxyPlugin;
import org.apache.shenyu.plugin.ai.proxy.enhanced.handler.AiProxyPluginHandler;
import org.apache.shenyu.plugin.ai.proxy.enhanced.protocol.AiProxyProtocol;
@@ -31,14 +32,18 @@ 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.subscriber.CommonAiProxyApiKeyDataSubscriber;
import org.apache.shenyu.plugin.ai.proxy.enhanced.transport.AiProxyTransport;
+import
org.apache.shenyu.plugin.ai.proxy.enhanced.transport.AiProxyTransportProperties;
import
org.apache.shenyu.plugin.ai.proxy.enhanced.transport.WebClientAiProxyTransport;
import org.apache.shenyu.plugin.api.ShenyuPlugin;
import org.apache.shenyu.sync.data.api.AiProxyApiKeyDataSubscriber;
import
org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
+import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
+import org.springframework.http.client.reactive.ReactorClientHttpConnector;
import org.springframework.web.reactive.function.client.WebClient;
+import reactor.netty.http.client.HttpClient;
import java.util.List;
@@ -55,9 +60,9 @@ public class AiProxyPluginConfiguration {
/**
* Ai proxy plugin.
*
- * @param aiProxyConfigService the aiProxyConfigService
- * @param aiProxyEngine the aiProxyEngine
- * @param aiProxyPluginHandler the aiProxyPluginHandler
+ * @param aiProxyConfigService the aiProxyConfigService
+ * @param aiProxyEngine the aiProxyEngine
+ * @param aiProxyPluginHandler the aiProxyPluginHandler
* @return the shenyu plugin
*/
@Bean
@@ -101,12 +106,12 @@ public class AiProxyPluginConfiguration {
*
* @param aiProxyProtocolFactory protocol factory
* @param aiProxyProviderFactory provider factory
- * @param aiProxyTransport transport
+ * @param aiProxyTransport transport
* @return the AI proxy engine
*/
@Bean
public AiProxyEngine aiProxyEngine(final AiProxyProtocolFactory
aiProxyProtocolFactory,
- final AiProxyProviderFactory aiProxyProviderFactory, final
AiProxyTransport aiProxyTransport) {
+ final AiProxyProviderFactory
aiProxyProviderFactory, final AiProxyTransport aiProxyTransport) {
return new AiProxyEngine(aiProxyProtocolFactory,
aiProxyProviderFactory, aiProxyTransport);
}
@@ -165,16 +170,19 @@ public class AiProxyPluginConfiguration {
/**
* Ai proxy transport.
*
+ * @param properties AI proxy transport properties
* @return the AI proxy transport
*/
@Bean
@ConditionalOnMissingBean(AiProxyTransport.class)
- public AiProxyTransport aiProxyTransport() {
+ public AiProxyTransport aiProxyTransport(final AiProxyTransportProperties
properties) {
final WebClient webClient = WebClient.builder()
+ .clientConnector(new
ReactorClientHttpConnector(createHttpClient(properties)))
.build();
- return new WebClientAiProxyTransport(webClient);
+ return new WebClientAiProxyTransport(webClient, properties);
}
+
/**
* Create the DeepSeek compatible provider profile.
*
@@ -215,4 +223,26 @@ public class AiProxyPluginConfiguration {
return new OpenAiCompatible("Moonshot");
}
+ /**
+ * AI proxy transport properties.
+ *
+ * @return AI proxy transport properties
+ */
+ @Bean
+ @ConfigurationProperties(prefix = "shenyu.plugins.ai.proxy.transport")
+ public AiProxyTransportProperties aiProxyTransportProperties() {
+ return new AiProxyTransportProperties();
+ }
+
+ HttpClient createHttpClient(final AiProxyTransportProperties properties) {
+ return HttpClient.create()
+ .option(
+ ChannelOption.CONNECT_TIMEOUT_MILLIS,
+ connectTimeoutMillis(properties)
+ );
+ }
+
+ int connectTimeoutMillis(final AiProxyTransportProperties properties) {
+ return Math.toIntExact(properties.getConnectTimeout().toMillis());
+ }
}
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 80aad054f2..9db0b79b02 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
@@ -24,20 +24,29 @@ 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.transport.AiProxyTransport;
+import
org.apache.shenyu.plugin.ai.proxy.enhanced.transport.AiProxyTransportProperties;
import org.junit.jupiter.api.Test;
+import org.springframework.boot.autoconfigure.AutoConfigurations;
+import
org.springframework.boot.autoconfigure.context.ConfigurationPropertiesAutoConfiguration;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import reactor.core.publisher.Mono;
+import java.time.Duration;
import java.util.List;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertSame;
/**
- * Auto-configuration for the migrated execution path and provider profiles.
+ * Auto-configuration for the migrated execution path, provider profiles and
transport properties.
*/
class AiProxyPluginConfigurationTest {
+ private final ApplicationContextRunner contextRunner = new
ApplicationContextRunner()
+ .withConfiguration(AutoConfigurations.of(
+ ConfigurationPropertiesAutoConfiguration.class,
+ AiProxyPluginConfiguration.class));
+
@Test
void createsTheMigratedPluginAndAllExistingProviderIdentifiers() {
new
ApplicationContextRunner().withUserConfiguration(AiProxyPluginConfiguration.class).run(context
-> {
@@ -62,4 +71,66 @@ class AiProxyPluginConfigurationTest {
new
ApplicationContextRunner().withUserConfiguration(AiProxyPluginConfiguration.class).withPropertyValues("shenyu.plugins.ai.proxy.enabled=false")
.run(context ->
org.assertj.core.api.Assertions.assertThat(context).doesNotHaveBean(AiProxyPlugin.class).doesNotHaveBean(AiProxyEngine.class));
}
+
+ @Test
+ void shouldUseDefaultConnectTimeout() {
+ contextRunner.run(context -> {
+ AiProxyTransportProperties properties =
context.getBean(AiProxyTransportProperties.class);
+ assertEquals(Duration.ofSeconds(30),
properties.getConnectTimeout());
+ assertEquals(Duration.ofSeconds(30),
properties.getResponseHeaderTimeout());
+ assertEquals(Duration.ofSeconds(60),
properties.getStreamingIdleTimeout());
+ assertEquals(Duration.ofMinutes(10),
properties.getTotalStreamDuration());
+ assertEquals(0, properties.getMaxResponseBytes());
+ });
+ }
+
+ @Test
+ void shouldBindConfiguredConnectTimeout() {
+ contextRunner
+ .withPropertyValues(
+ "shenyu.plugins.ai.proxy.transport.connect-timeout=5s",
+
"shenyu.plugins.ai.proxy.transport.response-header-timeout=7s",
+
"shenyu.plugins.ai.proxy.transport.streaming-idle-timeout=11s",
+
"shenyu.plugins.ai.proxy.transport.total-stream-duration=12m",
+
"shenyu.plugins.ai.proxy.transport.max-response-bytes=1024")
+ .run(context -> {
+ AiProxyTransportProperties properties =
context.getBean(AiProxyTransportProperties.class);
+ assertEquals(Duration.ofSeconds(5),
properties.getConnectTimeout());
+ assertEquals(Duration.ofSeconds(7),
properties.getResponseHeaderTimeout());
+ assertEquals(Duration.ofSeconds(11),
properties.getStreamingIdleTimeout());
+ assertEquals(Duration.ofMinutes(12),
properties.getTotalStreamDuration());
+ assertEquals(1024, properties.getMaxResponseBytes());
+ });
+ }
+
+ @Test
+ void shouldConvertConnectTimeoutToNettyMilliseconds() {
+ final AiProxyTransportProperties properties = new
AiProxyTransportProperties();
+ properties.setConnectTimeout(Duration.ofSeconds(5));
+
+ final int timeout = new AiProxyPluginConfiguration()
+ .connectTimeoutMillis(properties);
+
+ assertEquals(5000, timeout);
+ }
+
+ @Test
+ void shouldRejectNonPositiveTransportTimeouts() {
+ for (String property : List.of(
+ "connect-timeout",
+ "response-header-timeout",
+ "streaming-idle-timeout",
+ "total-stream-duration")) {
+ contextRunner
+ .withPropertyValues("shenyu.plugins.ai.proxy.transport." +
property + "=0s")
+ .run(context ->
org.assertj.core.api.Assertions.assertThat(context).hasFailed());
+ }
+ }
+
+ @Test
+ void shouldRejectNegativeMaximumResponseBytes() {
+ contextRunner
+
.withPropertyValues("shenyu.plugins.ai.proxy.transport.max-response-bytes=-1")
+ .run(context ->
org.assertj.core.api.Assertions.assertThat(context).hasFailed());
+ }
}
