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 4f58c24cae [ISSUE #6642] Capture flushed MCP response bodies. (#7396)
4f58c24cae is described below
commit 4f58c24caeb7427912b5af96b1730580181b7810
Author: Jerry聊AI <[email protected]>
AuthorDate: Sun Oct 11 23:20:34 2026 +0800
[ISSUE #6642] Capture flushed MCP response bodies. (#7396)
---
.../response/ShenyuMcpResponseDecorator.java | 65 ++++++++----
.../response/ShenyuMcpResponseDecoratorTest.java | 113 ++++++++++++++++++++-
2 files changed, 160 insertions(+), 18 deletions(-)
diff --git
a/shenyu-plugin/shenyu-plugin-mcp-server/src/main/java/org/apache/shenyu/plugin/mcp/server/response/ShenyuMcpResponseDecorator.java
b/shenyu-plugin/shenyu-plugin-mcp-server/src/main/java/org/apache/shenyu/plugin/mcp/server/response/ShenyuMcpResponseDecorator.java
index 6b59d8b42d..5499105b7f 100644
---
a/shenyu-plugin/shenyu-plugin-mcp-server/src/main/java/org/apache/shenyu/plugin/mcp/server/response/ShenyuMcpResponseDecorator.java
+++
b/shenyu-plugin/shenyu-plugin-mcp-server/src/main/java/org/apache/shenyu/plugin/mcp/server/response/ShenyuMcpResponseDecorator.java
@@ -27,6 +27,8 @@ import
org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
+import java.io.ByteArrayOutputStream;
+import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.util.Objects;
import java.util.concurrent.CompletableFuture;
@@ -35,7 +37,9 @@ public class ShenyuMcpResponseDecorator extends
ServerHttpResponseDecorator {
private static final Logger LOG =
LoggerFactory.getLogger(ShenyuMcpResponseDecorator.class);
- private final StringBuilder body = new StringBuilder();
+ private static final int MAX_CAPTURE_BYTES = 1024 * 1024;
+
+ private final ByteArrayOutputStream body = new ByteArrayOutputStream();
private final CompletableFuture<String> future;
@@ -43,6 +47,8 @@ public class ShenyuMcpResponseDecorator extends
ServerHttpResponseDecorator {
private boolean isFirstChunk = true;
+ private volatile boolean captureLimitExceeded;
+
private final JsonObject responseTemplate;
public ShenyuMcpResponseDecorator(final ServerHttpResponse response, final
String sessionId,
@@ -56,25 +62,17 @@ public class ShenyuMcpResponseDecorator extends
ServerHttpResponseDecorator {
@Override
public Mono<Void> writeWith(final Publisher<? extends DataBuffer> body) {
LOG.debug("Writing response data for session: {}", sessionId);
- return super.writeWith(Flux.from(body).doOnNext(buffer -> {
- byte[] bytes = new byte[buffer.readableByteCount()];
- buffer.read(bytes);
- String chunk = new String(bytes, StandardCharsets.UTF_8);
- if (isFirstChunk) {
- LOG.debug("First response chunk received for session: {}",
sessionId);
- isFirstChunk = false;
- }
- LOG.debug("Received response chunk for session {}, length: {}",
sessionId, chunk.length());
- synchronized (this.body) {
- this.body.append(chunk);
- }
- }).doOnComplete(() -> completeFuture()));
+ return super.writeWith(Flux.from(body).doOnNext(this::appendChunk))
+ .doOnSuccess(unused -> completeFuture());
}
@Override
public Mono<Void> writeAndFlushWith(final Publisher<? extends Publisher<?
extends DataBuffer>> body) {
LOG.debug("Writing and flushing response data for session: {}",
sessionId);
- return super.writeAndFlushWith(body);
+ final Flux<Publisher<? extends DataBuffer>> capturedBody =
Flux.from(body)
+ .map(inner -> Flux.from(inner).doOnNext(this::appendChunk));
+ return super.writeAndFlushWith(capturedBody)
+ .doOnSuccess(unused -> completeFuture());
}
@Override
@@ -85,10 +83,15 @@ public class ShenyuMcpResponseDecorator extends
ServerHttpResponseDecorator {
}
private void completeFuture() {
- String responseBody;
+ if (captureLimitExceeded) {
+ future.completeExceptionally(new IllegalStateException("MCP
response exceeds the 1 MiB capture limit"));
+ return;
+ }
+ byte[] responseBytes;
synchronized (this.body) {
- responseBody = this.body.toString();
+ responseBytes = this.body.toByteArray();
}
+ String responseBody = new String(responseBytes,
StandardCharsets.UTF_8);
LOG.debug("Final response body length for session {}: {}", sessionId,
responseBody.length());
if (!future.isDone()) {
synchronized (future) {
@@ -99,6 +102,34 @@ public class ShenyuMcpResponseDecorator extends
ServerHttpResponseDecorator {
}
}
+ private void appendChunk(final DataBuffer buffer) {
+ final int chunkLength = buffer.readableByteCount();
+ if (isFirstChunk) {
+ LOG.debug("First response chunk received for session: {}",
sessionId);
+ isFirstChunk = false;
+ }
+ LOG.debug("Received response chunk for session {}, length: {}",
sessionId, chunkLength);
+ synchronized (this.body) {
+ if (captureLimitExceeded) {
+ return;
+ }
+ try (DataBuffer.ByteBufferIterator iterator =
buffer.readableByteBuffers()) {
+ while (iterator.hasNext()) {
+ final ByteBuffer byteBuffer =
iterator.next().asReadOnlyBuffer();
+ final int remainingCapacity = MAX_CAPTURE_BYTES -
this.body.size();
+ if (byteBuffer.remaining() > remainingCapacity) {
+ captureLimitExceeded = true;
+ LOG.warn("Response body for session {} exceeds the 1
MiB capture limit", sessionId);
+ return;
+ }
+ final byte[] bytes = new byte[byteBuffer.remaining()];
+ byteBuffer.get(bytes);
+ this.body.write(bytes, 0, bytes.length);
+ }
+ }
+ }
+ }
+
private String applyResponseTemplate(final String responseBody) {
if (Objects.isNull(responseTemplate)) {
return responseBody;
diff --git
a/shenyu-plugin/shenyu-plugin-mcp-server/src/test/java/org/apache/shenyu/plugin/mcp/server/response/ShenyuMcpResponseDecoratorTest.java
b/shenyu-plugin/shenyu-plugin-mcp-server/src/test/java/org/apache/shenyu/plugin/mcp/server/response/ShenyuMcpResponseDecoratorTest.java
index 94f4269f71..d7b58f5251 100644
---
a/shenyu-plugin/shenyu-plugin-mcp-server/src/test/java/org/apache/shenyu/plugin/mcp/server/response/ShenyuMcpResponseDecoratorTest.java
+++
b/shenyu-plugin/shenyu-plugin-mcp-server/src/test/java/org/apache/shenyu/plugin/mcp/server/response/ShenyuMcpResponseDecoratorTest.java
@@ -17,6 +17,7 @@
package org.apache.shenyu.plugin.mcp.server.response;
+import org.reactivestreams.Publisher;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.Mock;
@@ -27,11 +28,16 @@ import
org.springframework.http.server.reactive.ServerHttpResponse;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
+import java.io.ByteArrayOutputStream;
import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.List;
import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.when;
@@ -48,7 +54,12 @@ class ShenyuMcpResponseDecoratorTest {
@Test
void testWriteWithCompletesFutureWithAllChunks() throws Exception {
- when(delegate.writeWith(any())).thenAnswer(invocation ->
Flux.from(invocation.getArgument(0)).then());
+ final List<String> forwardedChunks = new ArrayList<>();
+ when(delegate.writeWith(any())).thenAnswer(invocation -> Flux
+ .from(invocation.<Publisher<? extends
DataBuffer>>getArgument(0))
+ .map(this::readBuffer)
+ .doOnNext(forwardedChunks::add)
+ .then());
final CompletableFuture<String> future = new CompletableFuture<>();
final ShenyuMcpResponseDecorator decorator =
@@ -57,6 +68,90 @@ class ShenyuMcpResponseDecoratorTest {
decorator.writeWith(Flux.just(buffer("part-1,"),
buffer("part-2"))).block();
assertEquals("part-1,part-2", future.get(5, TimeUnit.SECONDS));
+ assertEquals(List.of("part-1,", "part-2"), forwardedChunks);
+ }
+
+ @Test
+ void testWriteWithPreservesUtf8CharacterSplitAcrossBuffers() throws
Exception {
+ final ByteArrayOutputStream forwardedBody = new
ByteArrayOutputStream();
+ when(delegate.writeWith(any())).thenAnswer(invocation -> Flux
+ .from(invocation.<Publisher<? extends
DataBuffer>>getArgument(0))
+ .doOnNext(buffer -> appendBytes(forwardedBody, buffer))
+ .then());
+
+ final CompletableFuture<String> future = new CompletableFuture<>();
+ final ShenyuMcpResponseDecorator decorator =
+ new ShenyuMcpResponseDecorator(delegate, "session-1", future,
null);
+ final byte[] utf8Character = "é".getBytes(StandardCharsets.UTF_8);
+
+ decorator.writeWith(Flux.just(buffer(new byte[]{utf8Character[0]}),
buffer(new byte[]{utf8Character[1]}))).block();
+
+ assertEquals("é", future.get(5, TimeUnit.SECONDS));
+ assertEquals("é", new String(forwardedBody.toByteArray(),
StandardCharsets.UTF_8));
+ }
+
+ @Test
+ void testWriteAndFlushWithCompletesFutureAndPreservesFlushes() throws
Exception {
+ final List<List<String>> forwardedFlushes = new ArrayList<>();
+ when(delegate.writeAndFlushWith(any())).thenAnswer(invocation -> Flux
+ .from(invocation.<Publisher<? extends Publisher<? extends
DataBuffer>>>getArgument(0))
+ .concatMap(flush -> Flux.from(flush)
+ .map(this::readBuffer)
+ .collectList()
+ .doOnNext(forwardedFlushes::add))
+ .then());
+
+ final CompletableFuture<String> future = new CompletableFuture<>();
+ final ShenyuMcpResponseDecorator decorator =
+ new ShenyuMcpResponseDecorator(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));
+ assertEquals(List.of(List.of("part-1,"), List.of("part-2")),
forwardedFlushes);
+ }
+
+ @Test
+ void testWriteAndFlushWithPreservesUtf8CharacterSplitAcrossFlushes()
throws Exception {
+ final ByteArrayOutputStream forwardedBody = new
ByteArrayOutputStream();
+ when(delegate.writeAndFlushWith(any())).thenAnswer(invocation -> Flux
+ .from(invocation.<Publisher<? extends Publisher<? extends
DataBuffer>>>getArgument(0))
+ .concatMap(flush -> Flux.from(flush).doOnNext(buffer ->
appendBytes(forwardedBody, buffer)).then())
+ .then());
+
+ final CompletableFuture<String> future = new CompletableFuture<>();
+ final ShenyuMcpResponseDecorator decorator =
+ new ShenyuMcpResponseDecorator(delegate, "session-1", future,
null);
+ final byte[] utf8Character = "é".getBytes(StandardCharsets.UTF_8);
+
+ decorator.writeAndFlushWith(Flux.just(
+ Flux.just(buffer(new byte[]{utf8Character[0]})),
+ Flux.just(buffer(new byte[]{utf8Character[1]})))).block();
+
+ assertEquals("é", future.get(5, TimeUnit.SECONDS));
+ assertEquals("é", new String(forwardedBody.toByteArray(),
StandardCharsets.UTF_8));
+ }
+
+ @Test
+ void testWriteWithRejectsCaptureLargerThanLimitButForwardsWholeBody()
throws Exception {
+ final ByteArrayOutputStream forwardedBody = new
ByteArrayOutputStream();
+ when(delegate.writeWith(any())).thenAnswer(invocation -> Flux
+ .from(invocation.<Publisher<? extends
DataBuffer>>getArgument(0))
+ .doOnNext(buffer -> appendBytes(forwardedBody, buffer))
+ .then());
+
+ final CompletableFuture<String> future = new CompletableFuture<>();
+ final ShenyuMcpResponseDecorator decorator =
+ new ShenyuMcpResponseDecorator(delegate, "session-1", future,
null);
+ final byte[] response = new byte[1024 * 1024 + 1];
+
+ decorator.writeWith(Flux.just(buffer(response))).block();
+
+ final ExecutionException exception =
assertThrows(ExecutionException.class, future::get);
+ assertEquals("MCP response exceeds the 1 MiB capture limit",
exception.getCause().getMessage());
+ assertEquals(response.length, forwardedBody.size());
}
@Test
@@ -75,4 +170,20 @@ class ShenyuMcpResponseDecoratorTest {
private DataBuffer buffer(final String content) {
return bufferFactory.wrap(content.getBytes(StandardCharsets.UTF_8));
}
+
+ private DataBuffer buffer(final byte[] content) {
+ return bufferFactory.wrap(content);
+ }
+
+ private void appendBytes(final ByteArrayOutputStream output, final
DataBuffer dataBuffer) {
+ final byte[] bytes = new byte[dataBuffer.readableByteCount()];
+ dataBuffer.read(bytes);
+ output.write(bytes, 0, bytes.length);
+ }
+
+ private String readBuffer(final DataBuffer dataBuffer) {
+ final byte[] bytes = new byte[dataBuffer.readableByteCount()];
+ dataBuffer.read(bytes);
+ return new String(bytes, StandardCharsets.UTF_8);
+ }
}
