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 410d552a6f99..6b59d8b42d8b 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 @@ -68,15 +68,7 @@ public Mono writeWith(final Publisher body) { synchronized (this.body) { this.body.append(chunk); } - // Complete future early for efficiency, but safely check if already done - if (!future.isDone()) { - synchronized (future) { - if (!future.isDone()) { - future.complete(applyResponseTemplate(this.body.toString())); - } - } - } - })); + }).doOnComplete(() -> completeFuture())); } @Override @@ -88,6 +80,11 @@ public Mono writeAndFlushWith(final Publisher setComplete() { LOG.debug("Response completed for session: {}", sessionId); + completeFuture(); + return super.setComplete(); + } + + private void completeFuture() { String responseBody; synchronized (this.body) { responseBody = this.body.toString(); @@ -100,7 +97,6 @@ public Mono setComplete() { } } } - return super.setComplete(); } private String applyResponseTemplate(final String 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 new file mode 100644 index 000000000000..94f4269f7179 --- /dev/null +++ b/shenyu-plugin/shenyu-plugin-mcp-server/src/test/java/org/apache/shenyu/plugin/mcp/server/response/ShenyuMcpResponseDecoratorTest.java @@ -0,0 +1,78 @@ +/* + * 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.mcp.server.response; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +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.http.server.reactive.ServerHttpResponse; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; + +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.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.when; + +/** + * Test case for {@link ShenyuMcpResponseDecorator}. + */ +@ExtendWith(MockitoExtension.class) +class ShenyuMcpResponseDecoratorTest { + + @Mock + private ServerHttpResponse delegate; + + private final DefaultDataBufferFactory bufferFactory = new DefaultDataBufferFactory(); + + @Test + void testWriteWithCompletesFutureWithAllChunks() throws Exception { + when(delegate.writeWith(any())).thenAnswer(invocation -> Flux.from(invocation.getArgument(0)).then()); + + final CompletableFuture future = new CompletableFuture<>(); + final ShenyuMcpResponseDecorator decorator = + new ShenyuMcpResponseDecorator(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)); + } + + @Test + void testSetCompleteCompletesFutureWithAccumulatedBody() { + when(delegate.setComplete()).thenReturn(Mono.empty()); + + final CompletableFuture future = new CompletableFuture<>(); + final ShenyuMcpResponseDecorator decorator = + new ShenyuMcpResponseDecorator(delegate, "session-1", future, null); + + decorator.setComplete().block(); + + assertEquals("", future.getNow("")); + } + + private DataBuffer buffer(final String content) { + return bufferFactory.wrap(content.getBytes(StandardCharsets.UTF_8)); + } +}