feat: improve thread safety and resource management in WebMVC streamable transport

- Add ReentrantLock synchronization to prevent race conditions in concurrent sseBuilder usage
- Add volatile closed flag to track session state and prevent operations on closed sessions
- Move SSE onComplete callback to properly close listening stream when connection completes
- Refactor closeGracefully to delegate to close() method for consistent cleanup
- Update class documentation to reflect thread-safety guarantees

This ensures safe concurrent access to SSE builder and proper cleanup of streaming resources.

Signed-off-by: Christian Tzolov <[email protected]>

Author: tzolov

mcp-spring/mcp-spring-webmvc/src/main/java/io/modelcontextprotocol/server/transport/WebMvcStreamableServerTransportProvider.java

@@ -7,6 +7,7 @@
 import java.io.IOException;
 import java.time.Duration;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.Function;
 
 import com.fasterxml.jackson.core.type.TypeReference;
@@ -230,9 +231,6 @@ private ServerResponse handleGet(ServerRequest request) {
 
 		try {
 			return ServerResponse.sse(sseBuilder -> {
-				sseBuilder.onComplete(() -> {
-					logger.debug("SSE connection completed for session: {}", sessionId);
-				});
 				sseBuilder.onTimeout(() -> {
 					logger.debug("SSE connection timed out for session: {}", sessionId);
 				});
@@ -264,8 +262,11 @@ private ServerResponse handleGet(ServerRequest request) {
 					// Establish new listening stream
 					McpStreamableServerSession.McpStreamableServerSessionStream listeningStream = session
 						.listeningStream(sessionTransport);
-					// Note: WebMVC SSE doesn't have onCancel, cleanup will happen in
-					// onComplete/onTimeout
+
+					sseBuilder.onComplete(() -> {
+						logger.debug("SSE connection completed for session: {}", sessionId);
+						listeningStream.close();
+					});
 				}
 			}, Duration.ZERO);
 		}
@@ -415,13 +416,22 @@ private ServerResponse handleDelete(ServerRequest request) {
 	/**
 	 * Implementation of McpStreamableServerTransport for WebMVC SSE sessions. This class
 	 * handles the transport-level communication for a specific client session.
+	 *
+	 * <p>
+	 * This class is thread-safe and uses a ReentrantLock to synchronize access to the
+	 * underlying SSE builder to prevent race conditions when multiple threads attempt to
+	 * send messages concurrently.
 	 */
 	private class WebMvcStreamableMcpSessionTransport implements McpStreamableServerTransport {
 
 		private final String sessionId;
 
 		private final SseBuilder sseBuilder;
 
+		private final ReentrantLock lock = new ReentrantLock();
+
+		private volatile boolean closed = false;
+
 		/**
 		 * Creates a new session transport with the specified ID and SSE builder.
 		 * @param sessionId The unique identifier for this session
@@ -453,7 +463,18 @@ public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message) {
 		@Override
 		public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message, String messageId) {
 			return Mono.fromRunnable(() -> {
+				if (this.closed) {
+					logger.debug("Attempted to send message to closed session: {}", this.sessionId);
+					return;
+				}
+
+				this.lock.lock();
 				try {
+					if (this.closed) {
+						logger.debug("Session {} was closed during message send attempt", this.sessionId);
+						return;
+					}
+
 					String jsonText = objectMapper.writeValueAsString(message);
 					this.sseBuilder.id(messageId != null ? messageId : this.sessionId)
 						.event(MESSAGE_EVENT_TYPE)
@@ -462,7 +483,16 @@ public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message, String messageId
 				}
 				catch (Exception e) {
 					logger.error("Failed to send message to session {}: {}", this.sessionId, e.getMessage());
-					this.sseBuilder.error(e);
+					try {
+						this.sseBuilder.error(e);
+					}
+					catch (Exception errorException) {
+						logger.error("Failed to send error to SSE builder for session {}: {}", this.sessionId,
+								errorException.getMessage());
+					}
+				}
+				finally {
+					this.lock.unlock();
 				}
 			});
 		}
@@ -486,14 +516,7 @@ public <T> T unmarshalFrom(Object data, TypeReference<T> typeRef) {
 		@Override
 		public Mono<Void> closeGracefully() {
 			return Mono.fromRunnable(() -> {
-				logger.debug("Closing streamable session transport: {}", this.sessionId);
-				try {
-					this.sseBuilder.complete();
-					logger.debug("Successfully completed SSE builder for session {}", this.sessionId);
-				}
-				catch (Exception e) {
-					logger.warn("Failed to complete SSE builder for session {}: {}", this.sessionId, e.getMessage());
-				}
+				WebMvcStreamableMcpSessionTransport.this.close();
 			});
 		}
 
@@ -502,13 +525,24 @@ public Mono<Void> closeGracefully() {
 		 */
 		@Override
 		public void close() {
+			this.lock.lock();
 			try {
+				if (this.closed) {
+					logger.debug("Session transport {} already closed", this.sessionId);
+					return;
+				}
+
+				this.closed = true;
+
 				this.sseBuilder.complete();
 				logger.debug("Successfully completed SSE builder for session {}", sessionId);
 			}
 			catch (Exception e) {
 				logger.warn("Failed to complete SSE builder for session {}: {}", sessionId, e.getMessage());
 			}
+			finally {
+				this.lock.unlock();
+			}
 		}
 
 	}