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 903363f9bf Isolate pending requests in MCP Streamable HTTP sessions
(#7638)
903363f9bf is described below
commit 903363f9bf3c7aa92d22acd97a0c66fe208bb7d5
Author: BobSong <[email protected]>
AuthorDate: Sat Oct 17 20:47:41 2026 +0800
Isolate pending requests in MCP Streamable HTTP sessions (#7638)
* fix(mcp): isolate pending Streamable HTTP request lifecycles
* style(mcp): group Streamable HTTP transport test imports
---------
Co-authored-by: BobSong-dev <[email protected]>
---
...henyuStreamableHttpServerTransportProvider.java | 151 +++++++---
...uStreamableHttpServerTransportProviderTest.java | 331 +++++++++++++++++++--
2 files changed, 420 insertions(+), 62 deletions(-)
diff --git
a/shenyu-plugin/shenyu-plugin-mcp-server/src/main/java/org/apache/shenyu/plugin/mcp/server/transport/ShenyuStreamableHttpServerTransportProvider.java
b/shenyu-plugin/shenyu-plugin-mcp-server/src/main/java/org/apache/shenyu/plugin/mcp/server/transport/ShenyuStreamableHttpServerTransportProvider.java
index ba356c2a14..b0bd801dbd 100644
---
a/shenyu-plugin/shenyu-plugin-mcp-server/src/main/java/org/apache/shenyu/plugin/mcp/server/transport/ShenyuStreamableHttpServerTransportProvider.java
+++
b/shenyu-plugin/shenyu-plugin-mcp-server/src/main/java/org/apache/shenyu/plugin/mcp/server/transport/ShenyuStreamableHttpServerTransportProvider.java
@@ -40,6 +40,7 @@ import
org.springframework.web.reactive.function.server.ServerResponse;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
+import reactor.core.publisher.Sinks;
import java.io.IOException;
import java.util.LinkedHashSet;
@@ -74,6 +75,8 @@ public class ShenyuStreamableHttpServerTransportProvider
implements McpServerTra
*/
private static final String JSONRPC_VERSION = "2.0";
+ private static final Object REQUEST_RESPONSE_CONTEXT_KEY = new Object();
+
/**
* MCP initialize method name.
*/
@@ -551,18 +554,27 @@ public class ShenyuStreamableHttpServerTransportProvider
implements McpServerTra
// Let MCP framework handle the message - framework will send response
through transport
final McpTransportContext transportContext =
McpTransportContext.create(
Map.of(McpSessionHelper.SHENYU_EXCHANGE_CONTEXT_KEY,
exchange));
- return Mono.defer(() -> session.handle(message))
- // The exchange belongs to this JSON-RPC request, not to the
session.
- .contextWrite(context -> context.put(McpTransportContext.KEY,
transportContext))
- .cast(Object.class)
- .doOnSuccess(result -> LOGGER.debug("Successfully processed
message for session: {}", sessionId))
- .then(waitForTransportResponse(transport, sessionId,
messageId))
- .doFinally(signal -> {
- // Release only this request's response on completion,
error or cancellation.
- if (Objects.nonNull(transport)) {
- transport.resetCapturedMessage(messageId);
- }
- })
+ return Mono.defer(() -> {
+ if (Objects.isNull(transport)) {
+ return Mono.error(new IllegalStateException("MCP session is
closed"));
+ }
+ final boolean requestMessage = message instanceof
McpSchema.JSONRPCRequest;
+ final PendingResponse pending = requestMessage &&
Objects.nonNull(messageId) ? transport.registerRequest(messageId) : null;
+ if (requestMessage && Objects.isNull(pending)) {
+ return
Mono.just(createMessageHandlingResult(HttpStatus.BAD_REQUEST.value(),
+ createJsonRpcError(messageId, -32600, "Invalid or
duplicate pending request id"), sessionId));
+ }
+ Mono<MessageHandlingResult> result = Mono.defer(() ->
session.handle(message))
+ .contextWrite(context -> Objects.isNull(pending) ?
context.put(McpTransportContext.KEY, transportContext)
+ : context.put(McpTransportContext.KEY,
transportContext).put(REQUEST_RESPONSE_CONTEXT_KEY, pending))
+ .then(waitForTransportResponse(transport, sessionId,
messageId, pending));
+ if (Objects.nonNull(pending)) {
+ final Mono<MessageHandlingResult> closed =
pending.closed.asMono().then(Mono.error(new IllegalStateException("MCP session
is closed")));
+ result = Mono.firstWithSignal(closed, result)
+ .doFinally(signal -> transport.finishRequest(pending));
+ }
+ return result;
+ })
.onErrorResume(error -> {
LOGGER.error("Error processing message for session {}:
{}", sessionId, error.getMessage(), error);
final Object errorResponse = createJsonRpcError(messageId,
-32603,
@@ -809,8 +821,7 @@ public class ShenyuStreamableHttpServerTransportProvider
implements McpServerTra
final McpServerSession removedSession = sessions.remove(sessionId);
final StreamableHttpSessionTransport removedTransport =
sessionTransports.remove(sessionId);
if (Objects.nonNull(removedTransport)) {
- removedTransport.closed = true;
- removedTransport.resetCapturedMessage();
+ removedTransport.closeRequests();
}
ShenyuMcpExchangeHolder.remove(sessionId);
if (Objects.nonNull(removedSession) ||
Objects.nonNull(removedTransport)) {
@@ -836,14 +847,17 @@ public class ShenyuStreamableHttpServerTransportProvider
implements McpServerTra
* @param transport the session transport to check for responses
* @param sessionId the session identifier for logging
* @param messageId the original message ID for correlation
+ * @param pending the response owned by this request, or null for
non-request messages
* @return a Mono containing the message handling result
*/
private Mono<MessageHandlingResult> waitForTransportResponse(final
StreamableHttpSessionTransport transport,
final String
sessionId,
- final Object
messageId) {
+ final Object
messageId,
+ final
PendingResponse pending) {
return Mono.fromCallable(() -> {
- if (Objects.nonNull(transport) &&
transport.isResponseReady(messageId)) {
- final McpSchema.JSONRPCMessage sentMessage =
transport.getResponse(messageId);
+ final McpSchema.JSONRPCMessage sentMessage =
Objects.nonNull(pending)
+ ? transport.getRequestResponse(pending) :
transport.getResponse(messageId);
+ if (Objects.nonNull(sentMessage)) {
LOGGER.debug("Retrieved captured response from transport for
session: {}", sessionId);
return createMessageHandlingResult(200, sentMessage,
sessionId);
} else {
@@ -1047,12 +1061,21 @@ public class
ShenyuStreamableHttpServerTransportProvider implements McpServerTra
}
}
- /**
- * Session transport implementation for Streamable HTTP with proper
lifecycle management.
- * This transport handles the communication between the MCP framework and
the Streamable HTTP
- * protocol. It captures responses from the MCP session and makes them
available for
- * HTTP response correlation.
- */
+ /** Response and close signal belonging to one subscribed HTTP request,
never just to an RPC id. */
+ private static final class PendingResponse {
+
+ private final Object messageId;
+
+ private final Sinks.Empty<Void> closed = Sinks.empty();
+
+ private McpSchema.JSONRPCMessage response;
+
+ private PendingResponse(final Object messageId) {
+ this.messageId = messageId;
+ }
+ }
+
+ /** Session transport with request-owned response capture and
cancellation. */
private class StreamableHttpSessionTransport implements McpServerTransport
{
private String sessionId;
@@ -1063,6 +1086,8 @@ public class ShenyuStreamableHttpServerTransportProvider
implements McpServerTra
private final ConcurrentHashMap<Object, McpSchema.JSONRPCMessage>
responses = new ConcurrentHashMap<>();
+ private final Map<Object, PendingResponse> pendingResponses = new
ConcurrentHashMap<>();
+
private volatile boolean responseReady;
/**
@@ -1138,17 +1163,65 @@ public class
ShenyuStreamableHttpServerTransportProvider implements McpServerTra
return responses.get(messageId);
}
+ private synchronized PendingResponse registerRequest(final Object
messageId) {
+ if (closed) {
+ throw new IllegalStateException("MCP session is closed");
+ }
+ if (pendingResponses.containsKey(messageId)) {
+ return null;
+ }
+ final PendingResponse pending = new PendingResponse(messageId);
+ pendingResponses.put(messageId, pending);
+ return pending;
+ }
+
+ private synchronized McpSchema.JSONRPCMessage getRequestResponse(final
PendingResponse pending) {
+ if (closed || pendingResponses.get(pending.messageId) != pending) {
+ throw new IllegalStateException("MCP request is no longer
active");
+ }
+ return pending.response;
+ }
+
+ private synchronized void finishRequest(final PendingResponse pending)
{
+ pendingResponses.remove(pending.messageId, pending);
+ pending.response = null;
+ }
+
+ private void closeRequests() {
+ final List<PendingResponse> pending;
+ synchronized (this) {
+ closed = true;
+ pending = List.copyOf(pendingResponses.values());
+ pendingResponses.clear();
+ pending.forEach(response -> response.response = null);
+ resetCapturedMessage();
+ }
+ // Signal outside the lock: cancellation can execute arbitrary
handler cleanup.
+ pending.forEach(response -> response.closed.tryEmitEmpty());
+ }
+
@Override
public Mono<Void> sendMessage(final McpSchema.JSONRPCMessage message) {
- if (!closed) {
- this.lastSentMessage = message;
- if (message instanceof McpSchema.JSONRPCResponse response &&
Objects.nonNull(response.id())) {
- responses.put(response.id(), message);
+ return Mono.deferContextual(context -> {
+ synchronized (this) {
+ if (closed) {
+ return Mono.empty();
+ }
+ final PendingResponse pending =
context.getOrDefault(REQUEST_RESPONSE_CONTEXT_KEY, null);
+ if (Objects.nonNull(pending)) {
+ if (pendingResponses.get(pending.messageId) == pending
&& message instanceof McpSchema.JSONRPCResponse response
+ && Objects.equals(pending.messageId,
response.id())) {
+ pending.response = message;
+ }
+ } else if (message instanceof McpSchema.JSONRPCResponse
response && "__backend_init".equals(response.id())) {
+ // Only the internal handshake uses the session-level
compatibility slots.
+ this.lastSentMessage = message;
+ responses.put(response.id(), message);
+ this.responseReady = true;
+ }
}
- this.responseReady = true;
- LOGGER.debug("Captured response message for session: {}",
sessionId);
- }
- return Mono.empty();
+ return Mono.empty();
+ });
}
@Override
@@ -1158,22 +1231,12 @@ public class
ShenyuStreamableHttpServerTransportProvider implements McpServerTra
@Override
public Mono<Void> closeGracefully() {
- return Mono.fromRunnable(() -> {
- if (!closed) {
- closed = true;
- removeSession(sessionId);
- LOGGER.debug("Session transport closed gracefully: {}",
sessionId);
- }
- });
+ return Mono.fromRunnable(() -> removeSession(sessionId));
}
@Override
public void close() {
- if (!closed) {
- closed = true;
- removeSession(sessionId);
- LOGGER.debug("Session transport closed immediately: {}",
sessionId);
- }
+ removeSession(sessionId);
}
/**
@@ -1181,7 +1244,7 @@ public class ShenyuStreamableHttpServerTransportProvider
implements McpServerTra
* business request is not confused with the internal initialize
handshake that
* is executed during short-reconnect or temporary session creation.
*/
- public void resetCapturedMessage() {
+ public synchronized void resetCapturedMessage() {
this.responses.clear();
this.lastSentMessage = null;
this.responseReady = false;
diff --git
a/shenyu-plugin/shenyu-plugin-mcp-server/src/test/java/org/apache/shenyu/plugin/mcp/server/transport/ShenyuStreamableHttpServerTransportProviderTest.java
b/shenyu-plugin/shenyu-plugin-mcp-server/src/test/java/org/apache/shenyu/plugin/mcp/server/transport/ShenyuStreamableHttpServerTransportProviderTest.java
index e73ccae0dc..0ac453ac0a 100644
---
a/shenyu-plugin/shenyu-plugin-mcp-server/src/test/java/org/apache/shenyu/plugin/mcp/server/transport/ShenyuStreamableHttpServerTransportProviderTest.java
+++
b/shenyu-plugin/shenyu-plugin-mcp-server/src/test/java/org/apache/shenyu/plugin/mcp/server/transport/ShenyuStreamableHttpServerTransportProviderTest.java
@@ -17,40 +17,56 @@
package org.apache.shenyu.plugin.mcp.server.transport;
+import java.lang.reflect.Field;
+import java.lang.reflect.Method;
+import java.time.Duration;
+import java.util.Collections;
+import java.util.List;
+import java.util.Locale;
+import java.util.Map;
+import java.util.Objects;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+
import com.fasterxml.jackson.databind.ObjectMapper;
+import io.modelcontextprotocol.server.McpRequestHandler;
import io.modelcontextprotocol.spec.McpSchema;
import io.modelcontextprotocol.spec.McpServerSession;
import io.modelcontextprotocol.spec.McpServerTransport;
-import io.modelcontextprotocol.server.McpRequestHandler;
-import org.apache.shenyu.plugin.mcp.server.session.McpSessionHelper;
-import org.springframework.web.server.ServerWebExchange;
-import reactor.core.publisher.Flux;
-import reactor.core.Disposable;
-import java.util.concurrent.atomic.AtomicBoolean;
import org.apache.shenyu.plugin.mcp.server.holder.ShenyuMcpExchangeHolder;
+import org.apache.shenyu.plugin.mcp.server.session.McpSessionHelper;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.Test;
import org.springframework.http.HttpStatus;
+import org.springframework.http.MediaType;
+import org.springframework.http.ResponseEntity;
import org.springframework.http.codec.HttpMessageWriter;
+import org.springframework.http.server.reactive.ReactorHttpHandlerAdapter;
import org.springframework.mock.http.server.reactive.MockServerHttpRequest;
import org.springframework.mock.http.server.reactive.MockServerHttpResponse;
import org.springframework.mock.web.server.MockServerWebExchange;
+import org.springframework.web.reactive.function.client.WebClient;
import org.springframework.web.reactive.function.server.HandlerStrategies;
+import org.springframework.web.reactive.function.server.RequestPredicates;
+import org.springframework.web.reactive.function.server.RouterFunctions;
import org.springframework.web.reactive.function.server.ServerRequest;
import org.springframework.web.reactive.function.server.ServerResponse;
import org.springframework.web.reactive.result.view.ViewResolver;
+import org.springframework.web.server.ServerWebExchange;
+import reactor.core.Disposable;
+import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
+import reactor.core.publisher.Sinks;
+import reactor.netty.http.server.HttpServer;
+import reactor.netty.resources.LoopResources;
import reactor.test.StepVerifier;
-
-import java.lang.reflect.Field;
-import java.lang.reflect.Method;
-import java.time.Duration;
-import java.util.Collections;
-import java.util.List;
-import java.util.Locale;
-import java.util.Map;
-import java.util.Objects;
-import java.util.UUID;
+import reactor.util.context.Context;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNotEquals;
@@ -393,10 +409,288 @@ class ShenyuStreamableHttpServerTransportProviderTest {
}
}
+ @Test
+ void testCompletedRequestReleasesLastResponse() throws Exception {
+ final ShenyuStreamableHttpServerTransportProvider provider =
providerWithRealSessions();
+ final String sessionId = initialize(provider);
+ try {
+ performRequest(provider, postRequest(TOOLS_LIST_REQUEST_BODY,
sessionId));
+ final Object transport = readMap(provider,
"sessionTransports").get(sessionId);
+ final Method lookup =
transport.getClass().getDeclaredMethod("getLastSentMessage");
+ lookup.setAccessible(true);
+ assertNull(lookup.invoke(transport), "A completed request must not
retain its result in the session slot");
+ } finally {
+ provider.removeSession(sessionId);
+ }
+ }
+
+ @Test
+ void testDuplicatePendingIdDoesNotInvokeSecondHandler() {
+ final AtomicInteger calls = new AtomicInteger();
+ final ShenyuStreamableHttpServerTransportProvider provider =
providerWithHandler((exchange, params) ->
+ calls.incrementAndGet() == 1 ? Mono.never() :
Mono.just(Map.of("marker", "second")));
+ final String sessionId = initialize(provider);
+ final Disposable pending =
provider.handleUnifiedEndpoint(createRequest(postRequest(TOOLS_LIST_REQUEST_BODY,
sessionId))).subscribe();
+ try {
+ assertEquals(1, calls.get());
+ final MockServerHttpResponse duplicate = performRequest(provider,
postRequest(TOOLS_LIST_REQUEST_BODY, sessionId));
+ assertEquals(HttpStatus.BAD_REQUEST, duplicate.getStatusCode());
+ assertEquals(1, calls.get(), "Duplicate in-flight IDs must be
rejected before a second handler executes");
+ } finally {
+ pending.dispose();
+ provider.removeSession(sessionId);
+ }
+ }
+
+ @Test
+ void testSessionRemovalCancelsPendingHandler() {
+ final AtomicBoolean cancelled = new AtomicBoolean();
+ final AtomicReference<ServerResponse> result = new AtomicReference<>();
+ final ShenyuStreamableHttpServerTransportProvider provider =
providerWithHandler((exchange, params) ->
+ Mono.<Map<String, Object>>never().doOnCancel(() ->
cancelled.set(true)));
+ final String sessionId = initialize(provider);
+ final Disposable pending =
provider.handleUnifiedEndpoint(createRequest(postRequest(TOOLS_LIST_REQUEST_BODY,
sessionId))).subscribe(result::set);
+ try {
+ provider.removeSession(sessionId);
+ assertTrue(cancelled.get(), "Closing a session must cancel its
owned handler subscription");
+ assertNotNull(result.get(), "A closed session must terminate the
waiting HTTP request");
+ assertEquals(HttpStatus.INTERNAL_SERVER_ERROR,
result.get().statusCode());
+ } finally {
+ pending.dispose();
+ provider.removeSession(sessionId);
+ }
+ }
+
+ @Test
+ void testLateResponseAfterCancellationIsDiscarded() throws Exception {
+ final AtomicReference<Context> requestContext = new
AtomicReference<>();
+ final AtomicReference<McpServerTransport> transport = new
AtomicReference<>();
+ final ShenyuStreamableHttpServerTransportProvider provider =
providerWithHandler((exchange, params) ->
+ Mono.deferContextual(context -> {
+ requestContext.set(Context.of(context));
+ return Mono.never();
+ }), transport);
+ final String sessionId = initialize(provider);
+ final Disposable pending =
provider.handleUnifiedEndpoint(createRequest(postRequest(TOOLS_LIST_REQUEST_BODY,
sessionId))).subscribe();
+ try {
+ assertNotNull(requestContext.get());
+ pending.dispose();
+ final McpSchema.JSONRPCResponse late = new
McpSchema.JSONRPCResponse("2.0", "req-1", Map.of("private", "late-result"),
null);
+
transport.get().sendMessage(late).contextWrite(requestContext.get()).block(Duration.ofSeconds(3));
+ final Field responses =
transport.get().getClass().getDeclaredField("responses");
+ responses.setAccessible(true);
+ assertTrue(((Map<?, ?>) responses.get(transport.get())).isEmpty(),
"Late responses must not repopulate cancelled request state");
+ } finally {
+ pending.dispose();
+ provider.removeSession(sessionId);
+ }
+ }
+
+ @Test
+ void testCancelledOwnerCannotOverwriteSameIdSuccessor() {
+ final AtomicReference<Context> oldContext = new AtomicReference<>();
+ final AtomicReference<McpServerTransport> transport = new
AtomicReference<>();
+ final AtomicInteger calls = new AtomicInteger();
+ final Sinks.One<Map<String, Object>> freshResult = Sinks.one();
+ final ShenyuStreamableHttpServerTransportProvider provider =
providerWithHandler((exchange, params) -> Mono.deferContextual(context -> {
+ if (calls.incrementAndGet() == 1) {
+ oldContext.set(Context.of(context));
+ return Mono.never();
+ }
+ return freshResult.asMono();
+ }), transport);
+ final String sessionId = initialize(provider);
+ final Disposable first =
provider.handleUnifiedEndpoint(createRequest(postRequest(TOOLS_LIST_REQUEST_BODY,
sessionId))).subscribe();
+ first.dispose();
+ final AtomicReference<ServerResponse> result = new AtomicReference<>();
+ final Disposable second =
provider.handleUnifiedEndpoint(createRequest(postRequest(TOOLS_LIST_REQUEST_BODY,
sessionId))).subscribe(result::set);
+ try {
+ assertEquals(2, calls.get());
+ transport.get().sendMessage(new McpSchema.JSONRPCResponse("2.0",
"req-1", Map.of("marker", "old-private"), null))
+
.contextWrite(oldContext.get()).block(Duration.ofSeconds(3));
+ assertNull(result.get());
+ freshResult.emitValue(Map.of("marker", "fresh"),
Sinks.EmitFailureHandler.FAIL_FAST);
+ assertNotNull(result.get());
+ final MockServerWebExchange output =
MockServerWebExchange.from(MockServerHttpRequest.get("/result"));
+ result.get().writeTo(output,
RESPONSE_CONTEXT).block(Duration.ofSeconds(3));
+ final String body =
output.getResponse().getBodyAsString().block(Duration.ofSeconds(3));
+ assertTrue(body.contains("\"marker\":\"fresh\""), body);
+ assertTrue(!body.contains("old-private"), body);
+ } finally {
+ first.dispose();
+ second.dispose();
+ provider.removeSession(sessionId);
+ }
+ }
+
+ @Test
+ void testSameIdInDifferentSessionsKeepsPrivateResults() {
+ final ShenyuStreamableHttpServerTransportProvider provider =
providerWithHandler((exchange, params) -> {
+ final ServerWebExchange request = (ServerWebExchange)
exchange.transportContext().get(McpSessionHelper.SHENYU_EXCHANGE_CONTEXT_KEY);
+ return
Mono.delay(Duration.ofMillis(10)).thenReturn(Map.of("marker",
request.getRequest().getHeaders().getFirst("X-Test-Request")));
+ });
+ final List<String> sessions = List.of(initialize(provider),
initialize(provider));
+ try {
+ StepVerifier.create(Flux.range(0, 2).flatMap(index -> {
+ final String marker = "private-" + index;
+ final MockServerWebExchange exchange =
MockServerWebExchange.from(MockServerHttpRequest.post("/mcp/streamablehttp")
+ .header("Content-Type",
"application/json").header(SESSION_ID_HEADER, sessions.get(index))
+ .header("X-Test-Request",
marker).body(TOOLS_LIST_REQUEST_BODY));
+ return
provider.handleUnifiedEndpoint(ServerRequest.create(exchange,
HandlerStrategies.withDefaults().messageReaders()))
+ .flatMap(response -> response.writeTo(exchange,
RESPONSE_CONTEXT))
+ .then(Mono.defer(() ->
exchange.getResponse().getBodyAsString()))
+ .doOnNext(body ->
assertTrue(body.contains("\"marker\":\"" + marker + "\""), body));
+ })).expectNextCount(2).verifyComplete();
+ } finally {
+ sessions.forEach(provider::removeSession);
+ }
+ }
+
+ @Test
+ void testClosingOneSessionLeavesAnotherRequestActive() {
+ final AtomicBoolean firstCancelled = new AtomicBoolean();
+ final AtomicBoolean secondCancelled = new AtomicBoolean();
+ final Sinks.One<Map<String, Object>> secondResult = Sinks.one();
+ final ShenyuStreamableHttpServerTransportProvider provider =
providerWithHandler((exchange, params) -> {
+ final ServerWebExchange request = (ServerWebExchange)
exchange.transportContext().get(McpSessionHelper.SHENYU_EXCHANGE_CONTEXT_KEY);
+ return
"close".equals(request.getRequest().getHeaders().getFirst("X-Test-Request"))
+ ? Mono.<Map<String, Object>>never().doOnCancel(() ->
firstCancelled.set(true))
+ : secondResult.asMono().doOnCancel(() ->
secondCancelled.set(true));
+ });
+ final String firstSession = initialize(provider);
+ final String secondSession = initialize(provider);
+ final Disposable first =
provider.handleUnifiedEndpoint(createRequest(MockServerHttpRequest.post("/mcp/streamablehttp")
+ .header("Content-Type",
"application/json").header(SESSION_ID_HEADER, firstSession)
+ .header("X-Test-Request",
"close").body(TOOLS_LIST_REQUEST_BODY))).subscribe();
+ final AtomicReference<ServerResponse> result = new AtomicReference<>();
+ final Disposable second =
provider.handleUnifiedEndpoint(createRequest(postRequest(TOOLS_LIST_REQUEST_BODY,
secondSession))).subscribe(result::set);
+ try {
+ provider.removeSession(firstSession);
+ assertTrue(firstCancelled.get());
+ assertTrue(!secondCancelled.get());
+ assertNull(result.get());
+ secondResult.emitValue(Map.of("tools", List.of()),
Sinks.EmitFailureHandler.FAIL_FAST);
+ assertNotNull(result.get());
+ assertEquals(HttpStatus.OK, result.get().statusCode());
+ } finally {
+ first.dispose();
+ second.dispose();
+ provider.removeSession(firstSession);
+ provider.removeSession(secondSession);
+ }
+ }
+
+ @Test
+ void testCallerDeadlineReleasesPendingRequest() throws Exception {
+ final AtomicBoolean cancelled = new AtomicBoolean();
+ final ShenyuStreamableHttpServerTransportProvider provider =
providerWithHandler((exchange, params) ->
+ Mono.<Map<String, Object>>never().doOnCancel(() ->
cancelled.set(true)));
+ final String sessionId = initialize(provider);
+ try {
+ StepVerifier.withVirtualTime(() ->
provider.handleUnifiedEndpoint(createRequest(postRequest(TOOLS_LIST_REQUEST_BODY,
sessionId)))
+ .timeout(Duration.ofSeconds(1)))
+
.thenAwait(Duration.ofSeconds(1)).expectError(TimeoutException.class).verify(Duration.ofSeconds(3));
+ assertTrue(cancelled.get());
+ final Object transport = readMap(provider,
"sessionTransports").get(sessionId);
+ final Field pending =
transport.getClass().getDeclaredField("pendingResponses");
+ pending.setAccessible(true);
+ assertTrue(((Map<?, ?>) pending.get(transport)).isEmpty());
+ } finally {
+ provider.removeSession(sessionId);
+ }
+ }
+
+ @Test
+ void testGracefulShutdownCancelsPendingRequest() throws Exception {
+ final AtomicBoolean cancelled = new AtomicBoolean();
+ final ShenyuStreamableHttpServerTransportProvider provider =
providerWithHandler((exchange, params) ->
+ Mono.<Map<String, Object>>never().doOnCancel(() ->
cancelled.set(true)));
+ final String sessionId = initialize(provider);
+ final Disposable pending =
provider.handleUnifiedEndpoint(createRequest(postRequest(TOOLS_LIST_REQUEST_BODY,
sessionId))).subscribe();
+ try {
+ StepVerifier.create(provider.closeGracefully()).verifyComplete();
+ assertTrue(cancelled.get());
+ assertTrue(readMap(provider, "sessions").isEmpty());
+ assertTrue(readMap(provider, "sessionTransports").isEmpty());
+ assertNull(ShenyuMcpExchangeHolder.get(sessionId));
+ } finally {
+ pending.dispose();
+ provider.removeSession(sessionId);
+ }
+ }
+
+ @Test
+ void testRealHttpDuplicateAndDisconnectKeepPeerHealthy() throws Exception {
+ final CountDownLatch started = new CountDownLatch(1);
+ final CountDownLatch cancelled = new CountDownLatch(1);
+ final AtomicInteger calls = new AtomicInteger();
+ final ShenyuStreamableHttpServerTransportProvider provider =
providerWithHandler((exchange, params) -> {
+ calls.incrementAndGet();
+ final ServerWebExchange request = (ServerWebExchange)
exchange.transportContext().get(McpSessionHelper.SHENYU_EXCHANGE_CONTEXT_KEY);
+ final String marker =
request.getRequest().getHeaders().getFirst("X-Test-Request");
+ return "pending".equals(marker) ? Mono.<Map<String,
Object>>never().doOnSubscribe(subscription -> started.countDown())
+ .doOnCancel(cancelled::countDown) :
Mono.just(Map.of("marker", marker));
+ });
+ final LoopResources loops = LoopResources.create("mcp-lifecycle-test",
2, true);
+ final ReactorHttpHandlerAdapter adapter = new
ReactorHttpHandlerAdapter(RouterFunctions.toHttpHandler(
+
RouterFunctions.route(RequestPredicates.POST("/mcp/streamablehttp"),
provider::handleUnifiedEndpoint)));
+ final reactor.netty.DisposableServer server =
HttpServer.create().host("127.0.0.1").port(0).runOn(loops).handle(adapter).bindNow();
+ final WebClient client = WebClient.create("http://127.0.0.1:"; +
server.port());
+ CompletableFuture<ResponseEntity<String>> pending = null;
+ String sessionId = null;
+ try {
+ final ResponseEntity<String> initialized = networkPost(client,
INITIALIZE_REQUEST_BODY, "", "init").block(Duration.ofSeconds(10));
+ assertNotNull(initialized);
+ sessionId = initialized.getHeaders().getFirst(SESSION_ID_HEADER);
+ assertNotNull(sessionId);
+ assertEquals(HttpStatus.ACCEPTED, networkPost(client,
INITIALIZED_NOTIFICATION_BODY, sessionId, "init")
+ .block(Duration.ofSeconds(10)).getStatusCode());
+ pending = networkPost(client, TOOLS_LIST_REQUEST_BODY, sessionId,
"pending").toFuture();
+ assertTrue(started.await(10, TimeUnit.SECONDS));
+ final ResponseEntity<String> duplicate = networkPost(client,
TOOLS_LIST_REQUEST_BODY, sessionId, "duplicate").block(Duration.ofSeconds(10));
+ assertNotNull(duplicate);
+ assertEquals(HttpStatus.BAD_REQUEST, duplicate.getStatusCode());
+ assertEquals(1, calls.get());
+ final ResponseEntity<String> peer = networkPost(client,
TOOLS_LIST_REQUEST_BODY.replace("req-1", "req-2"), sessionId, "peer")
+ .block(Duration.ofSeconds(10));
+ assertNotNull(peer);
+ assertTrue(peer.getBody().contains("\"marker\":\"peer\""),
peer.getBody());
+ pending.cancel(true);
+ assertTrue(cancelled.await(10, TimeUnit.SECONDS), "Real HTTP
disconnect must cancel the owned SDK handler");
+ final ResponseEntity<String> next = networkPost(client,
TOOLS_LIST_REQUEST_BODY.replace("req-1", "req-3"), sessionId, "next")
+ .block(Duration.ofSeconds(10));
+ assertNotNull(next);
+ assertTrue(next.getBody().contains("\"marker\":\"next\""),
next.getBody());
+ assertEquals(3, calls.get());
+ } finally {
+ if (Objects.nonNull(pending)) {
+ pending.cancel(true);
+ }
+ if (Objects.nonNull(sessionId)) {
+ provider.removeSession(sessionId);
+ }
+ server.disposeNow();
+ loops.disposeLater().block(Duration.ofSeconds(10));
+ }
+ }
+
+ private Mono<ResponseEntity<String>> networkPost(final WebClient client,
final String body, final String sessionId, final String marker) {
+ return
client.post().uri("/mcp/streamablehttp").contentType(MediaType.APPLICATION_JSON)
+ .header(SESSION_ID_HEADER, sessionId).header("X-Test-Request",
marker).bodyValue(body)
+ .exchangeToMono(response -> response.toEntity(String.class));
+ }
+
private ShenyuStreamableHttpServerTransportProvider
providerWithHandler(final McpRequestHandler<Map<String, Object>> handler) {
+ return providerWithHandler(handler, new AtomicReference<>());
+ }
+
+ private ShenyuStreamableHttpServerTransportProvider
providerWithHandler(final McpRequestHandler<Map<String, Object>> handler,
+
final AtomicReference<McpServerTransport> capturedTransport) {
ShenyuStreamableHttpServerTransportProvider provider =
new ShenyuStreamableHttpServerTransportProvider(new
ObjectMapper(), "/mcp/streamablehttp");
- provider.setSessionFactory(transport -> new McpServerSession(
+ provider.setSessionFactory(transport -> {
+ capturedTransport.set(transport);
+ return new McpServerSession(
UUID.randomUUID().toString(),
Duration.ofSeconds(10),
transport,
@@ -406,7 +700,8 @@ class ShenyuStreamableHttpServerTransportProviderTest {
new McpSchema.Implementation("ShenyuMcpServer",
"1.0.0"),
"test")),
Map.<String, McpRequestHandler<?>>of("tools/list", handler),
- Map.of()));
+ Map.of());
+ });
return provider;
}
