From 9601a9318233c2573c9f9d3351cc1e3cb79a9227 Mon Sep 17 00:00:00 2001 From: lxq19991111 <15060002560@163.com> Date: Fri, 28 Aug 2026 09:32:43 +0800 Subject: [PATCH 1/2] Close WebMVC Streamable transports on disconnect Register SSE lifecycle cleanup so a client that disconnects without a clean DELETE no longer leaves the listening stream, replay stream, or streaming-response transport dangling. Route every lifecycle event through the same idempotent transport close() instead of a separate markClosed() flag, and route SSE write failures through the close path so a write failure closes only the current transport without removing the logical MCP session. Closes #6860 Signed-off-by: lxq19991111 <15060002560@163.com> --- ...bMvcStreamableServerTransportProvider.java | 76 +++++++++---------- ...treamableServerTransportProviderTests.java | 75 ++++++++++++++++++ 2 files changed, 112 insertions(+), 39 deletions(-) diff --git a/mcp/transport/mcp-spring-webmvc/src/main/java/org/springframework/ai/mcp/server/webmvc/transport/WebMvcStreamableServerTransportProvider.java b/mcp/transport/mcp-spring-webmvc/src/main/java/org/springframework/ai/mcp/server/webmvc/transport/WebMvcStreamableServerTransportProvider.java index e0ca528f05..e7478777a1 100644 --- a/mcp/transport/mcp-spring-webmvc/src/main/java/org/springframework/ai/mcp/server/webmvc/transport/WebMvcStreamableServerTransportProvider.java +++ b/mcp/transport/mcp-spring-webmvc/src/main/java/org/springframework/ai/mcp/server/webmvc/transport/WebMvcStreamableServerTransportProvider.java @@ -352,19 +352,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)) @@ -379,7 +374,7 @@ private ServerResponse handleGet(ServerRequest request) { if (logger.isErrorEnabled()) { logger.error("Failed to replay message: " + e.getMessage()); } - sseBuilder.error(e); + sessionTransport.close(); } }); } @@ -387,7 +382,7 @@ private ServerResponse handleGet(ServerRequest request) { if (logger.isErrorEnabled()) { logger.error("Failed to replay messages: " + e.getMessage()); } - sseBuilder.error(e); + sessionTransport.close(); } } else { @@ -395,12 +390,7 @@ private ServerResponse handleGet(ServerRequest request) { 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); } @@ -487,19 +477,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) @@ -510,7 +491,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); } @@ -690,6 +671,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(); } @@ -703,7 +709,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; @@ -774,15 +780,7 @@ public Mono 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()); - } - } + this.close(); } finally { this.lock.unlock(); diff --git a/mcp/transport/mcp-spring-webmvc/src/test/java/org/springframework/ai/mcp/server/webmvc/transport/WebMvcStreamableServerTransportProviderTests.java b/mcp/transport/mcp-spring-webmvc/src/test/java/org/springframework/ai/mcp/server/webmvc/transport/WebMvcStreamableServerTransportProviderTests.java index 92521eaeb6..a336427a63 100644 --- a/mcp/transport/mcp-spring-webmvc/src/test/java/org/springframework/ai/mcp/server/webmvc/transport/WebMvcStreamableServerTransportProviderTests.java +++ b/mcp/transport/mcp-spring-webmvc/src/test/java/org/springframework/ai/mcp/server/webmvc/transport/WebMvcStreamableServerTransportProviderTests.java @@ -16,22 +16,30 @@ 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 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; @@ -193,4 +201,71 @@ 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 completeCaptor = ArgumentCaptor.captor(); + verify(sseBuilder).onComplete(completeCaptor.capture()); + completeCaptor.getValue().run(); + + ArgumentCaptor timeoutCaptor = ArgumentCaptor.captor(); + verify(sseBuilder).onTimeout(timeoutCaptor.capture()); + timeoutCaptor.getValue().run(); + + ArgumentCaptor> 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()); + transport.sendMessage(message, "message-1").block(); + + verify(sseBuilder).complete(); + + // 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(); + } + } From 2a7c6e7b881a0e284fe936bf10c124521e57e42f Mon Sep 17 00:00:00 2001 From: lxq19991111 <15060002560@163.com> Date: Thu, 24 Sep 2026 21:15:00 +0800 Subject: [PATCH 2/2] Propagate WebMVC SSE write failures An SSE write failure only closed the current WebMVC session transport and then let the send complete normally, so callers kept writing to a stream that no longer leads anywhere. Surface the failure as a McpTransportException instead. The transport still closes only the current connection and the logical MCP session is preserved, so a client can still delete or reconnect with the same session id and Last-Event-ID after a transient write failure. This matches what modelcontextprotocol/java-sdk#1130 settled on upstream. Broadcasting is unaffected: notifyClients logs and swallows per-session failures, and the keep-alive scheduler completes ping failures, so propagating a single write failure cannot break other sessions. The write-failure test now asserts the propagated exception rather than verifying that the SseBuilder was completed on the mock. That assertion was a false positive: DefaultSseBuilder sets sendFailed on the first failed write, after which complete() returns immediately, so the mock verified a call that cannot happen on a real connection. Signed-off-by: lxq19991111 <15060002560@163.com> --- .../WebMvcStreamableServerTransportProvider.java | 14 ++++++++++++-- ...bMvcStreamableServerTransportProviderTests.java | 8 ++++++-- 2 files changed, 18 insertions(+), 4 deletions(-) diff --git a/mcp/transport/mcp-spring-webmvc/src/main/java/org/springframework/ai/mcp/server/webmvc/transport/WebMvcStreamableServerTransportProvider.java b/mcp/transport/mcp-spring-webmvc/src/main/java/org/springframework/ai/mcp/server/webmvc/transport/WebMvcStreamableServerTransportProvider.java index e7478777a1..258063e73e 100644 --- a/mcp/transport/mcp-spring-webmvc/src/main/java/org/springframework/ai/mcp/server/webmvc/transport/WebMvcStreamableServerTransportProvider.java +++ b/mcp/transport/mcp-spring-webmvc/src/main/java/org/springframework/ai/mcp/server/webmvc/transport/WebMvcStreamableServerTransportProvider.java @@ -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; @@ -735,7 +736,10 @@ class WebMvcStreamableMcpSessionTransport implements McpStreamableServerTranspor /** * 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 sendMessage(McpSchema.JSONRPCMessage message) { @@ -747,7 +751,10 @@ public Mono 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 sendMessage(McpSchema.JSONRPCMessage message, @Nullable String messageId) { @@ -780,7 +787,10 @@ public Mono sendMessage(McpSchema.JSONRPCMessage message, @Nullable String if (logger.isErrorEnabled()) { logger.error("Failed to send message to session " + this.sessionId + ": " + e.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(); diff --git a/mcp/transport/mcp-spring-webmvc/src/test/java/org/springframework/ai/mcp/server/webmvc/transport/WebMvcStreamableServerTransportProviderTests.java b/mcp/transport/mcp-spring-webmvc/src/test/java/org/springframework/ai/mcp/server/webmvc/transport/WebMvcStreamableServerTransportProviderTests.java index a336427a63..68efad78df 100644 --- a/mcp/transport/mcp-spring-webmvc/src/test/java/org/springframework/ai/mcp/server/webmvc/transport/WebMvcStreamableServerTransportProviderTests.java +++ b/mcp/transport/mcp-spring-webmvc/src/test/java/org/springframework/ai/mcp/server/webmvc/transport/WebMvcStreamableServerTransportProviderTests.java @@ -25,6 +25,7 @@ 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; @@ -258,9 +259,12 @@ void sseWriteFailureClosesOnlyCurrentTransportWithoutRemovingSession() throws Ex var transport = provider.createSessionTransport(sessionId, sseBuilder); McpSchema.JSONRPCMessage message = new McpSchema.JSONRPCNotification("2.0", "server/notification", Map.of()); - transport.sendMessage(message, "message-1").block(); - verify(sseBuilder).complete(); + // 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