This is an automated email from the ASF dual-hosted git repository.
Aias00 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 0ce19b07c8 fix(ai): transform flushed response bodies (#7123)
0ce19b07c8 is described below
commit 0ce19b07c8e5b672f9767dbec7e3931595b8b14a
Author: Liming Deng <[email protected]>
AuthorDate: Sun Oct 11 23:18:45 2026 +0800
fix(ai): transform flushed response bodies (#7123)
---
.../response/AiResponseTransformerPlugin.java | 13 +++++-
.../response/AiResponseTransformerPluginTest.java | 53 ++++++++++++++++++++++
2 files changed, 64 insertions(+), 2 deletions(-)
diff --git
a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-response-transformer/src/main/java/org/apache/shenyu/plugin/ai/transformer/response/AiResponseTransformerPlugin.java
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-response-transformer/src/main/java/org/apache/shenyu/plugin/ai/transformer/response/AiResponseTransformerPlugin.java
index c8db6f4419..71ccc28e6a 100644
---
a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-response-transformer/src/main/java/org/apache/shenyu/plugin/ai/transformer/response/AiResponseTransformerPlugin.java
+++
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-response-transformer/src/main/java/org/apache/shenyu/plugin/ai/transformer/response/AiResponseTransformerPlugin.java
@@ -44,6 +44,7 @@ import org.springframework.http.codec.HttpMessageReader;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.lang.NonNull;
import org.springframework.web.server.ServerWebExchange;
+import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import org.reactivestreams.Publisher;
@@ -67,6 +68,7 @@ import com.fasterxml.jackson.databind.node.ObjectNode;
/**
* this is ai response transformer plugin.
+ * The complete response is buffered before transformation; streaming and SSE
responses are not supported.
*/
public class AiResponseTransformerPlugin extends AbstractShenyuPlugin {
@@ -305,7 +307,7 @@ public class AiResponseTransformerPlugin extends
AbstractShenyuPlugin {
@NonNull
public Mono<Void> writeWith(@NonNull final Publisher<? extends
DataBuffer> body) {
final Mono<DataBuffer> dataBufferMono = DataBufferUtils.join(body);
- return dataBufferMono.flatMap(dataBuffer -> {
+ return dataBufferMono.map(dataBuffer -> {
byte[] bytes = new byte[dataBuffer.readableByteCount()];
dataBuffer.read(bytes);
DataBufferUtils.release(dataBuffer);
@@ -375,7 +377,14 @@ public class AiResponseTransformerPlugin extends
AbstractShenyuPlugin {
return
WebFluxResultUtils.result(this.exchange,
finalResponseBody.getBytes(StandardCharsets.UTF_8));
});
});
- });
+ }).defaultIfEmpty(Mono.defer(() -> super.writeWith(Flux.empty())))
+ .flatMap(response -> response);
+ }
+
+ @Override
+ @NonNull
+ public Mono<Void> writeAndFlushWith(@NonNull final Publisher<? extends
Publisher<? extends DataBuffer>> body) {
+ return writeWith(Flux.from(body).concatMap(Flux::from));
}
private String extractBodyFromAiResponse(final String aiResponse) {
diff --git
a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-response-transformer/src/test/java/org/apache/shenyu/plugin/ai/transformer/response/AiResponseTransformerPluginTest.java
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-response-transformer/src/test/java/org/apache/shenyu/plugin/ai/transformer/response/AiResponseTransformerPluginTest.java
index df27629feb..63f1a802e4 100644
---
a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-response-transformer/src/test/java/org/apache/shenyu/plugin/ai/transformer/response/AiResponseTransformerPluginTest.java
+++
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-response-transformer/src/test/java/org/apache/shenyu/plugin/ai/transformer/response/AiResponseTransformerPluginTest.java
@@ -33,6 +33,8 @@ import org.apache.shenyu.plugin.api.result.ShenyuResult;
import org.apache.shenyu.plugin.api.utils.SpringBeanUtils;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.Mock;
import org.mockito.MockedConstruction;
@@ -62,12 +64,14 @@ import com.fasterxml.jackson.databind.ObjectMapper;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.RETURNS_DEEP_STUBS;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.mockConstruction;
import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoInteractions;
import static org.mockito.Mockito.when;
import static org.mockito.Mockito.lenient;
@@ -180,6 +184,55 @@ class AiResponseTransformerPluginTest {
}
}
+ @ParameterizedTest
+ @ValueSource(ints = {0, 1, 2})
+ void testEmptyResponseIsCommitted(final int mode) {
+ MockServerHttpResponse response = (MockServerHttpResponse)
exchange.getResponse();
+ response.setStatusCode(HttpStatus.NO_CONTENT);
+ response.getHeaders().set("X-Original", "retained");
+ AiResponseTransformerTemplate template =
mock(AiResponseTransformerTemplate.class);
+ ChatClient client = mock(ChatClient.class);
+ AiResponseTransformerPlugin.AiResponseTransformerDecorator decorator =
+ new
AiResponseTransformerPlugin.AiResponseTransformerDecorator(exchange, template,
client);
+ Mono<Void> completion;
+ if (mode == 0) {
+ completion = decorator.writeWith(Flux.empty());
+ } else if (mode == 1) {
+ completion = decorator.writeAndFlushWith(Flux.empty());
+ } else {
+ completion = decorator.writeAndFlushWith(Flux.just(Flux.empty(),
Flux.empty()));
+ }
+ StepVerifier.create(completion).verifyComplete();
+ assertTrue(response.isCommitted());
+ assertEquals(HttpStatus.NO_CONTENT, response.getStatusCode());
+ assertEquals("retained", response.getHeaders().getFirst("X-Original"));
+ verifyNoInteractions(template, client);
+ }
+
+ @Test
+ void testWriteAndFlushWithTransformsResponse() {
+ MockServerHttpResponse response = (MockServerHttpResponse)
exchange.getResponse();
+ AiResponseTransformerTemplate template =
mock(AiResponseTransformerTemplate.class);
+
when(template.assembleMessage(exchange)).thenReturn(Mono.just("{\"response\":{\"body\":\"\"}}"));
+ ChatClient chatClient = mock(ChatClient.class, RETURNS_DEEP_STUBS);
+ when(chatClient.prompt().user(anyString()).stream().content())
+ .thenReturn(Flux.just("HTTP/1.1 200 OK\nContent-Type:
application/json\n\n{\"transformed\":true}"));
+ ConfigurableApplicationContext context =
mock(ConfigurableApplicationContext.class);
+ when(context.getBean(ShenyuResult.class)).thenReturn(new
DefaultShenyuResult());
+ SpringBeanUtils.getInstance().setApplicationContext(context);
+ AiResponseTransformerPlugin.AiResponseTransformerDecorator decorator =
+ new
AiResponseTransformerPlugin.AiResponseTransformerDecorator(exchange, template,
chatClient);
+ DataBuffer first =
response.bufferFactory().wrap("orig".getBytes(StandardCharsets.UTF_8));
+ DataBuffer second =
response.bufferFactory().wrap("inal".getBytes(StandardCharsets.UTF_8));
+
+
StepVerifier.create(decorator.writeAndFlushWith(Flux.just(Flux.just(first),
Flux.just(second))))
+ .verifyComplete();
+
+ StepVerifier.create(response.getBodyAsString())
+ .expectNext("{\"transformed\":true}")
+ .verifyComplete();
+ }
+
@Test
void testOriginalHeadersRemainWhenTransformedBodyIsInvalid() {
MockServerHttpResponse response = (MockServerHttpResponse)
exchange.getResponse();
