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 a5f412e6f7 fix: stream AI and MCP response chunks instead of buffering
via collectList (#6795) (#7484)
a5f412e6f7 is described below
commit a5f412e6f73d9dc919cb96575d67a65eec441ed5
Author: wy471x <[email protected]>
AuthorDate: Thu Oct 15 09:55:20 2026 +0800
fix: stream AI and MCP response chunks instead of buffering via collectList
(#6795) (#7484)
- Accumulate AI transformer streamed content into a single StringBuilder
instead of collecting every chunk into a List and joining afterwards.
- Make NonCommittingMcpResponseDecorator append data buffer chunks
incrementally, use concatMap to keep writeAndFlushWith inner publishers
ordered, and complete the response future once the stream completes.
- Add unit tests covering incremental accumulation, ordering, response
template application and error propagation.
---
.../request/AiRequestTransformerPlugin.java | 7 +-
.../response/AiResponseTransformerPlugin.java | 7 +-
.../NonCommittingMcpResponseDecorator.java | 70 +++++++--------
.../NonCommittingMcpResponseDecoratorTest.java | 99 +++++++++++++++++++++-
4 files changed, 137 insertions(+), 46 deletions(-)
diff --git
a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-request-transformer/src/main/java/org/apache/shenyu/plugin/ai/transformer/request/AiRequestTransformerPlugin.java
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-request-transformer/src/main/java/org/apache/shenyu/plugin/ai/transformer/request/AiRequestTransformerPlugin.java
index 198c2db3c9..97ad730b73 100644
---
a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-request-transformer/src/main/java/org/apache/shenyu/plugin/ai/transformer/request/AiRequestTransformerPlugin.java
+++
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-request-transformer/src/main/java/org/apache/shenyu/plugin/ai/transformer/request/AiRequestTransformerPlugin.java
@@ -59,7 +59,6 @@ import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.regex.Pattern;
-import java.util.stream.Collectors;
import java.util.stream.Stream;
/**
@@ -137,9 +136,9 @@ public class AiRequestTransformerPlugin extends
AbstractShenyuPlugin {
return aiRequestTransformerTemplate.assembleMessage()
.flatMap(message ->
finalClient.prompt().user(message).stream().content()
- .collectList()
- .map(list -> Objects.isNull(list) ? "" :
list.stream().filter(Objects::nonNull)
- .collect(Collectors.joining("")))
+ .filter(Objects::nonNull)
+ .collect(StringBuilder::new, StringBuilder::append)
+ .map(StringBuilder::toString)
.flatMap(aiResponse -> {
LOG.debug("Request rewritten to: {}", aiResponse);
return convertHeader(exchange, aiResponse)
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 48d79f853f..492a7bf4e3 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
@@ -59,7 +59,6 @@ import java.util.List;
import java.util.Objects;
import java.util.Optional;
import java.util.regex.Pattern;
-import java.util.stream.Collectors;
import java.util.stream.Stream;
import java.util.zip.GZIPInputStream;
@@ -356,9 +355,9 @@ public class AiResponseTransformerPlugin extends
AbstractShenyuPlugin {
final String finalMessage =
messageWithResponseBody;
return
chatClient.prompt().user(finalMessage).stream().content()
- .collectList()
- .map(list -> Objects.isNull(list) ? "" :
list.stream().filter(Objects::nonNull)
- .collect(Collectors.joining("")))
+ .filter(Objects::nonNull)
+ .collect(StringBuilder::new,
StringBuilder::append)
+ .map(StringBuilder::toString)
.flatMap(aiResponse -> {
HttpHeaders newHeaders =
extractHeadersFromAiResponse(aiResponse);
diff --git
a/shenyu-plugin/shenyu-plugin-mcp-server/src/main/java/org/apache/shenyu/plugin/mcp/server/response/NonCommittingMcpResponseDecorator.java
b/shenyu-plugin/shenyu-plugin-mcp-server/src/main/java/org/apache/shenyu/plugin/mcp/server/response/NonCommittingMcpResponseDecorator.java
index 23c759ac81..489ea64ca5 100644
---
a/shenyu-plugin/shenyu-plugin-mcp-server/src/main/java/org/apache/shenyu/plugin/mcp/server/response/NonCommittingMcpResponseDecorator.java
+++
b/shenyu-plugin/shenyu-plugin-mcp-server/src/main/java/org/apache/shenyu/plugin/mcp/server/response/NonCommittingMcpResponseDecorator.java
@@ -59,6 +59,8 @@ public class NonCommittingMcpResponseDecorator extends
ServerHttpResponseDecorat
private final JsonObject responseTemplate;
+ private final StringBuilder body = new StringBuilder();
+
/**
* Constructs a new non-committing MCP response decorator.
*
@@ -83,9 +85,9 @@ public class NonCommittingMcpResponseDecorator extends
ServerHttpResponseDecorat
LOG.debug("Processing writeWith for session: {}", sessionId);
return Flux.from(body)
- .collectList()
.doOnDiscard(DataBuffer.class, DataBufferUtils::release)
- .doOnNext(this::processResponseData)
+ .doOnNext(this::appendBuffer)
+ .doOnComplete(this::completeResponseFuture)
.then()
.doOnSuccess(aVoid -> LOG.debug("Successfully completed
writeWith for session: {}", sessionId))
.doOnError(error -> handleProcessingError("writeWith", error));
@@ -96,10 +98,10 @@ public class NonCommittingMcpResponseDecorator extends
ServerHttpResponseDecorat
LOG.debug("Processing writeAndFlushWith for session: {}", sessionId);
return Flux.from(body)
- .flatMap(Flux::from)
- .collectList()
+ .concatMap(Flux::from)
.doOnDiscard(DataBuffer.class, DataBufferUtils::release)
- .doOnNext(this::processResponseData)
+ .doOnNext(this::appendBuffer)
+ .doOnComplete(this::completeResponseFuture)
.then()
.doOnSuccess(aVoid -> LOG.debug("Successfully completed
writeAndFlushWith for session: {}", sessionId))
.doOnError(error -> handleProcessingError("writeAndFlushWith",
error));
@@ -116,18 +118,38 @@ public class NonCommittingMcpResponseDecorator extends
ServerHttpResponseDecorat
}
/**
- * Processes the collected response data buffers and completes the
response future.
- *
- * <p>Aggregates all data buffers into a single response string, applies
response template
- * transformations if configured, and completes the response future with
the processed result.
+ * Appends a data buffer chunk to the accumulated response body as it is
emitted,
+ * so that chunks are processed incrementally instead of being buffered as
a list first.
*
- * @param dataBuffers the collected response data buffers
+ * @param buffer the emitted data buffer chunk
*/
- private void processResponseData(final java.util.List<? extends
DataBuffer> dataBuffers) {
+ private void appendBuffer(final DataBuffer buffer) {
try {
- final String responseBody = aggregateDataBuffers(dataBuffers);
- LOG.debug("Captured response data for session {}, length: {}
chars", sessionId, responseBody.length());
+ final byte[] bytes = new byte[buffer.readableByteCount()];
+ buffer.read(bytes);
+ synchronized (body) {
+ body.append(new String(bytes, StandardCharsets.UTF_8));
+ }
+ } finally {
+ DataBufferUtils.release(buffer);
+ }
+ }
+ /**
+ * Completes the response future with the accumulated response data after
all chunks
+ * have been received.
+ *
+ * <p>Applies response template transformations if configured, and
completes the response
+ * future with the processed result.
+ */
+ private void completeResponseFuture() {
+ final String responseBody;
+ synchronized (body) {
+ responseBody = body.toString();
+ }
+ LOG.debug("Captured response data for session {}, length: {} chars",
sessionId, responseBody.length());
+
+ try {
final String processedResponse = processResponse(responseBody);
LOG.debug("Processed response for session {}, final length: {}
chars", sessionId, processedResponse.length());
@@ -140,28 +162,6 @@ public class NonCommittingMcpResponseDecorator extends
ServerHttpResponseDecorat
}
}
- /**
- * Aggregates multiple data buffers into a single response string.
- *
- * @param dataBuffers the list of data buffers to aggregate
- * @return the aggregated response string
- */
- private String aggregateDataBuffers(final java.util.List<? extends
DataBuffer> dataBuffers) {
- final StringBuilder responseBuilder = new StringBuilder();
-
- try {
- for (DataBuffer buffer : dataBuffers) {
- final byte[] bytes = new byte[buffer.readableByteCount()];
- buffer.read(bytes);
- responseBuilder.append(new String(bytes,
StandardCharsets.UTF_8));
- }
- } finally {
- dataBuffers.forEach(DataBufferUtils::release);
- }
-
- return responseBuilder.toString();
- }
-
/**
* Handles processing errors by logging and completing the future
exceptionally.
*
diff --git
a/shenyu-plugin/shenyu-plugin-mcp-server/src/test/java/org/apache/shenyu/plugin/mcp/server/response/NonCommittingMcpResponseDecoratorTest.java
b/shenyu-plugin/shenyu-plugin-mcp-server/src/test/java/org/apache/shenyu/plugin/mcp/server/response/NonCommittingMcpResponseDecoratorTest.java
index 5f1484a2ca..3287a7d49e 100644
---
a/shenyu-plugin/shenyu-plugin-mcp-server/src/test/java/org/apache/shenyu/plugin/mcp/server/response/NonCommittingMcpResponseDecoratorTest.java
+++
b/shenyu-plugin/shenyu-plugin-mcp-server/src/test/java/org/apache/shenyu/plugin/mcp/server/response/NonCommittingMcpResponseDecoratorTest.java
@@ -17,28 +17,117 @@
package org.apache.shenyu.plugin.mcp.server.response;
+import com.google.gson.JsonObject;
import io.netty.buffer.PooledByteBufAllocator;
import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.ValueSource;
+import org.mockito.Mock;
+import org.mockito.junit.jupiter.MockitoExtension;
+import org.springframework.core.io.buffer.DataBuffer;
+import org.springframework.core.io.buffer.DefaultDataBufferFactory;
import org.springframework.core.io.buffer.NettyDataBuffer;
import org.springframework.core.io.buffer.NettyDataBufferFactory;
+import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.mock.http.server.reactive.MockServerHttpResponse;
-import reactor.core.publisher.Mono;
import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
import reactor.test.StepVerifier;
import java.nio.charset.StandardCharsets;
import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.TimeUnit;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
/**
* Test cases for {@link NonCommittingMcpResponseDecorator}.
*/
+@ExtendWith(MockitoExtension.class)
public final class NonCommittingMcpResponseDecoratorTest {
+ @Mock
+ private ServerHttpResponse delegate;
+
+ private final DefaultDataBufferFactory bufferFactory = new
DefaultDataBufferFactory();
+
+ // ---------------------------------------------------------------
+ // Chunks are accumulated incrementally instead of via collectList
+ // ---------------------------------------------------------------
+
+ @Test
+ void testWriteWithAccumulatesChunksAndCompletesFuture() throws Exception {
+ final CompletableFuture<String> future = new CompletableFuture<>();
+ final NonCommittingMcpResponseDecorator decorator =
+ new NonCommittingMcpResponseDecorator(delegate, "session-1",
future, null);
+
+ decorator.writeWith(Flux.just(buffer("part-1,"),
buffer("part-2"))).block();
+
+ assertEquals("part-1,part-2", future.get(5, TimeUnit.SECONDS));
+ verify(delegate, never()).writeWith(any());
+ }
+
+ @Test
+ void testWriteAndFlushWithProcessesInnerPublishersInOrder() throws
Exception {
+ final CompletableFuture<String> future = new CompletableFuture<>();
+ final NonCommittingMcpResponseDecorator decorator =
+ new NonCommittingMcpResponseDecorator(delegate, "session-1",
future, null);
+
+ decorator.writeAndFlushWith(Flux.just(Flux.just(buffer("part-1,")),
Flux.just(buffer("part-2")))).block();
+
+ assertEquals("part-1,part-2", future.get(5, TimeUnit.SECONDS));
+ verify(delegate, never()).writeAndFlushWith(any());
+ }
+
+ @Test
+ void testWriteWithAppliesResponseTemplateOnAccumulatedBody() throws
Exception {
+ final JsonObject contentTemplate = new JsonObject();
+ contentTemplate.addProperty("type", "text");
+ contentTemplate.addProperty("text", "${result}");
+ final JsonObject responseTemplate = new JsonObject();
+ responseTemplate.add("content", contentTemplate);
+
+ final CompletableFuture<String> future = new CompletableFuture<>();
+ final NonCommittingMcpResponseDecorator decorator =
+ new NonCommittingMcpResponseDecorator(delegate, "session-1",
future, responseTemplate);
+
+ decorator.writeWith(Flux.just(buffer("{\"result\":"),
buffer("\"ok\"}"))).block();
+
+ assertEquals("{\"content\":{\"type\":\"text\",\"text\":\"ok\"}}",
future.get(5, TimeUnit.SECONDS));
+ }
+
+ @Test
+ void testWriteWithCompletesFutureExceptionallyOnError() {
+ final CompletableFuture<String> future = new CompletableFuture<>();
+ final NonCommittingMcpResponseDecorator decorator =
+ new NonCommittingMcpResponseDecorator(delegate, "session-1",
future, null);
+
+ StepVerifier.create(decorator.writeWith(Flux.error(new
IllegalStateException("boom"))))
+ .verifyError(IllegalStateException.class);
+
+ assertTrue(future.isCompletedExceptionally());
+ }
+
+ @Test
+ void testSetCompleteCompletesFutureWithEmptyBody() throws Exception {
+ final CompletableFuture<String> future = new CompletableFuture<>();
+ final NonCommittingMcpResponseDecorator decorator =
+ new NonCommittingMcpResponseDecorator(delegate, "session-1",
future, null);
+
+ StepVerifier.create(decorator.setComplete()).verifyComplete();
+
+ assertEquals("", future.get(5, TimeUnit.SECONDS));
+ }
+
+ // ---------------------------------------------
+ // Emitted buffers must always be released again
+ // ---------------------------------------------
+
@ParameterizedTest
@ValueSource(booleans = {false, true})
void releasesCollectedBufferOnUpstreamError(final boolean flush) {
@@ -72,8 +161,8 @@ public final class NonCommittingMcpResponseDecoratorTest {
@Test
public void testWriteWithReleasesCollectedBuffers() {
- NettyDataBufferFactory bufferFactory = new
NettyDataBufferFactory(PooledByteBufAllocator.DEFAULT);
- NettyDataBuffer buffer = (NettyDataBuffer)
bufferFactory.wrap("response".getBytes(StandardCharsets.UTF_8));
+ NettyDataBufferFactory nettyBufferFactory = new
NettyDataBufferFactory(PooledByteBufAllocator.DEFAULT);
+ NettyDataBuffer buffer = (NettyDataBuffer)
nettyBufferFactory.wrap("response".getBytes(StandardCharsets.UTF_8));
CompletableFuture<String> responseFuture = new CompletableFuture<>();
NonCommittingMcpResponseDecorator decorator = new
NonCommittingMcpResponseDecorator(
new MockServerHttpResponse(), "session", responseFuture, null);
@@ -83,4 +172,8 @@ public final class NonCommittingMcpResponseDecoratorTest {
assertEquals("response", responseFuture.join());
assertEquals(0, buffer.getNativeBuffer().refCnt());
}
+
+ private DataBuffer buffer(final String content) {
+ return bufferFactory.wrap(content.getBytes(StandardCharsets.UTF_8));
+ }
}
