Skip to content

Close WebMvcStreamableServerTransportProvider transports on disconnect and propagate write failures #6871

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
lxq19991111 wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Close WebMvcStreamableServerTransportProvider transports on disconnect and propagate write failures #6871

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9601a93
Close WebMVC Streamable transports on disconnect
lxq19991111 Aug 28, 2026
2a7c6e7
Propagate WebMVC SSE write failures
lxq19991111 Sep 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
90 changes: 49 additions & 41 deletions ...ringframework/ai/mcp/server/webmvc/transport/WebMvcStreamableServerTransportProvider.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@
import io.modelcontextprotocol.spec.McpStreamableServerSession;
import io.modelcontextprotocol.spec.McpStreamableServerTransport;
import io.modelcontextprotocol.spec.McpStreamableServerTransportProvider;
import io.modelcontextprotocol.spec.McpTransportException;
import io.modelcontextprotocol.spec.ProtocolVersions;
import io.modelcontextprotocol.util.Assert;
import io.modelcontextprotocol.util.KeepAliveScheduler;
Expand Down Expand Up @@ -352,19 +353,14 @@ private ServerResponse handleGet(ServerRequest request) {

try {
return ServerResponse.sse(sseBuilder -> {
sseBuilder.onTimeout(() -> {
if (logger.isDebugEnabled()) {
logger.debug("SSE connection timed out for session: " + sessionId);
}
});

WebMvcStreamableMcpSessionTransport sessionTransport = new WebMvcStreamableMcpSessionTransport(
sessionId, sseBuilder);
WebMvcStreamableMcpSessionTransport sessionTransport = createSessionTransport(sessionId, sseBuilder);

// Check if this is a replay request
if (!request.headers().header(HttpHeaders.LAST_EVENT_ID).isEmpty()) {
String lastId = request.headers().asHttpHeaders().getFirst(HttpHeaders.LAST_EVENT_ID);

registerSseLifecycle(sseBuilder, sessionId, sessionTransport::close);

try {
session.replay(lastId)
.contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext))
Expand All @@ -379,28 +375,23 @@ private ServerResponse handleGet(ServerRequest request) {
if (logger.isErrorEnabled()) {
logger.error("Failed to replay message: " + e.getMessage());
}
sseBuilder.error(e);
sessionTransport.close();
}
});
}
catch (Exception e) {
if (logger.isErrorEnabled()) {
logger.error("Failed to replay messages: " + e.getMessage());
}
sseBuilder.error(e);
sessionTransport.close();
}
}
else {
// Establish new listening stream
McpStreamableServerSession.McpStreamableServerSessionStream listeningStream = session
.listeningStream(sessionTransport);

sseBuilder.onComplete(() -> {
if (logger.isDebugEnabled()) {
logger.debug("SSE connection completed for session: " + sessionId);
}
listeningStream.close();
});
registerSseLifecycle(sseBuilder, sessionId, listeningStream::close);
}
}, Duration.ZERO);
}
Expand Down Expand Up @@ -487,19 +478,10 @@ else if (message instanceof McpSchema.JSONRPCNotification jsonrpcNotification) {
else if (message instanceof McpSchema.JSONRPCRequest jsonrpcRequest) {
// For streaming responses, we need to return SSE
return ServerResponse.sse(sseBuilder -> {
sseBuilder.onComplete(() -> {
if (logger.isDebugEnabled()) {
logger.debug("Request response stream completed for session: " + sessionId);
}
});
sseBuilder.onTimeout(() -> {
if (logger.isDebugEnabled()) {
logger.debug("Request response stream timed out for session: " + sessionId);
}
});
WebMvcStreamableMcpSessionTransport sessionTransport = createSessionTransport(sessionId,
sseBuilder);

WebMvcStreamableMcpSessionTransport sessionTransport = new WebMvcStreamableMcpSessionTransport(
sessionId, sseBuilder);
registerSseLifecycle(sseBuilder, sessionId, sessionTransport::close);

try {
session.responseStream(jsonrpcRequest, sessionTransport)
Expand All @@ -510,7 +492,7 @@ else if (message instanceof McpSchema.JSONRPCRequest jsonrpcRequest) {
if (logger.isErrorEnabled()) {
logger.error("Failed to handle request stream: " + e.getMessage());
}
sseBuilder.error(e);
sessionTransport.close();
}
}, Duration.ZERO);
}
Expand Down Expand Up @@ -690,6 +672,31 @@ private McpStreamableServerSession removeSession(String sessionId) {
return this.sessions.remove(sessionId);
}

static void registerSseLifecycle(SseBuilder sseBuilder, String sessionId, Runnable onClose) {
sseBuilder.onComplete(() -> {
if (logger.isDebugEnabled()) {
logger.debug("SSE connection completed for session: " + sessionId);
}
onClose.run();
});
sseBuilder.onTimeout(() -> {
if (logger.isDebugEnabled()) {
logger.debug("SSE connection timed out for session: " + sessionId);
}
onClose.run();
});
sseBuilder.onError(error -> {
if (logger.isDebugEnabled()) {
logger.debug("SSE connection errored for session: " + sessionId, error);
}
onClose.run();
});
}

WebMvcStreamableMcpSessionTransport createSessionTransport(String sessionId, SseBuilder sseBuilder) {
return new WebMvcStreamableMcpSessionTransport(sessionId, sseBuilder);
}

public static Builder builder() {
return new Builder();
}
Expand All @@ -703,7 +710,7 @@ public static Builder builder() {
* underlying SSE builder to prevent race conditions when multiple threads attempt to
* send messages concurrently.
*/
private class WebMvcStreamableMcpSessionTransport implements McpStreamableServerTransport {
class WebMvcStreamableMcpSessionTransport implements McpStreamableServerTransport {

private final String sessionId;

Expand All @@ -729,7 +736,10 @@ private class WebMvcStreamableMcpSessionTransport implements McpStreamableServer
/**
* Sends a JSON-RPC message to the client through the SSE connection.
* @param message The JSON-RPC message to send
* @return A Mono that completes when the message has been sent
* @return A Mono that completes when the message has been sent, or fails with an
* {@link McpTransportException} if the SSE write fails, in which case this
* transport is closed but the logical MCP session is preserved for later
* reconnects
*/
@Override
public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message) {
Expand All @@ -741,7 +751,10 @@ public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message) {
* specific message ID.
* @param message The JSON-RPC message to send
* @param messageId The message ID for SSE event identification
* @return A Mono that completes when the message has been sent
* @return A Mono that completes when the message has been sent, or fails with an
* {@link McpTransportException} if the SSE write fails, in which case this
* transport is closed but the logical MCP session is preserved for later
* reconnects
*/
@Override
public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message, @Nullable String messageId) {
Expand Down Expand Up @@ -774,15 +787,10 @@ public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message, @Nullable String
if (logger.isErrorEnabled()) {
logger.error("Failed to send message to session " + this.sessionId + ": " + e.getMessage());
}
try {
this.sseBuilder.error(e);
}
catch (Exception errorException) {
if (logger.isErrorEnabled()) {
logger.error("Failed to send error to SSE builder for session " + this.sessionId + ": "
+ errorException.getMessage());
}
}
// Close only this transport: the logical MCP session is preserved so
// that a client can still reconnect with the same session id.
this.close();
throw new McpTransportException("Failed to send message to session " + this.sessionId, e);
}
finally {
this.lock.unlock();
Expand Down
79 changes: 79 additions & 0 deletions ...ramework/ai/mcp/server/webmvc/transport/WebMvcStreamableServerTransportProviderTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,22 +16,31 @@

package org.springframework.ai.mcp.server.webmvc.transport;

import java.io.IOException;
import java.time.Duration;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Consumer;

import io.modelcontextprotocol.spec.HttpHeaders;
import io.modelcontextprotocol.spec.McpSchema;
import io.modelcontextprotocol.spec.McpStreamableServerSession;
import io.modelcontextprotocol.spec.McpTransportException;
import org.junit.jupiter.api.Test;
import org.mockito.ArgumentCaptor;
import reactor.core.publisher.Mono;

import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.test.web.reactive.server.WebTestClient;
import org.springframework.test.web.servlet.client.MockMvcWebTestClient;
import org.springframework.web.servlet.function.ServerResponse.SseBuilder;

import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.timeout;
import static org.mockito.Mockito.verify;
Expand Down Expand Up @@ -193,4 +202,74 @@ void initializeRequestIsRejectedWhenSessionIdHeaderIsPresent() {
.value(body -> assertThat(body).contains("Session already initialized"));
}

@Test
void sseLifecycleCallbacksCloseStreamOnCompleteTimeoutAndError() {
SseBuilder sseBuilder = mock(SseBuilder.class);
AtomicInteger closeCount = new AtomicInteger();

WebMvcStreamableServerTransportProvider.registerSseLifecycle(sseBuilder, "session-1",
closeCount::incrementAndGet);

ArgumentCaptor<Runnable> completeCaptor = ArgumentCaptor.captor();
verify(sseBuilder).onComplete(completeCaptor.capture());
completeCaptor.getValue().run();

ArgumentCaptor<Runnable> timeoutCaptor = ArgumentCaptor.captor();
verify(sseBuilder).onTimeout(timeoutCaptor.capture());
timeoutCaptor.getValue().run();

ArgumentCaptor<Consumer<Throwable>> errorCaptor = ArgumentCaptor.captor();
verify(sseBuilder).onError(errorCaptor.capture());
errorCaptor.getValue().accept(new RuntimeException("boom"));

assertThat(closeCount).hasValue(3);
}

@Test
void sseWriteFailureClosesOnlyCurrentTransportWithoutRemovingSession() throws Exception {
String sessionId = "session-send-failure";
McpStreamableServerSession session = mock(McpStreamableServerSession.class);
when(session.getId()).thenReturn(sessionId);
when(session.delete()).thenReturn(Mono.empty());

McpSchema.InitializeResult initResult = McpSchema.InitializeResult
.builder("2024-11-05", McpSchema.ServerCapabilities.builder().build(),
McpSchema.Implementation.builder("test-server", "1.0.0").build())
.build();

WebMvcStreamableServerTransportProvider provider = WebMvcStreamableServerTransportProvider.builder().build();
provider.setSessionFactory(initializeRequest -> new McpStreamableServerSession.McpStreamableServerSessionInit(
session, Mono.just(initResult)));

WebTestClient client = MockMvcWebTestClient.bindToRouterFunction(provider.getRouterFunction()).build();

client.post()
.uri("/mcp")
.contentType(MediaType.APPLICATION_JSON)
.accept(MediaType.APPLICATION_JSON, MediaType.TEXT_EVENT_STREAM)
.bodyValue(INITIALIZE_REQUEST)
.exchange()
.expectStatus()
.isOk();

SseBuilder sseBuilder = mock(SseBuilder.class);
when(sseBuilder.id(anyString())).thenReturn(sseBuilder);
when(sseBuilder.event(anyString())).thenReturn(sseBuilder);
doThrow(new IOException("broken pipe")).when(sseBuilder).data(any());

var transport = provider.createSessionTransport(sessionId, sseBuilder);
McpSchema.JSONRPCMessage message = new McpSchema.JSONRPCNotification("2.0", "server/notification", Map.of());

// The write failure must surface to the caller instead of completing normally,
// so that callers stop writing to a stream that no longer leads anywhere.
assertThatThrownBy(() -> transport.sendMessage(message, "message-1").block())
.isInstanceOf(McpTransportException.class)
.hasCauseInstanceOf(IOException.class);

// A DELETE with the same id still finds the session instead of returning 404,
// proving the write failure closed only the transport and did not remove the
// logical MCP session.
client.delete().uri("/mcp").header(HttpHeaders.MCP_SESSION_ID, sessionId).exchange().expectStatus().isOk();
}

}