feat: add thread safety to SSE operations in WebMvcSseServerTransportProvider

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

Author: tzolov

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

@@ -8,6 +8,7 @@
 import java.time.Duration;
 import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.locks.ReentrantLock;
 
 import com.fasterxml.jackson.core.type.TypeReference;
 import com.fasterxml.jackson.databind.ObjectMapper;
@@ -339,6 +340,12 @@ private class WebMvcMcpSessionTransport implements McpServerTransport {
 
 		private final SseBuilder sseBuilder;
 
+		/**
+		 * Lock to ensure thread-safe access to the SSE builder when sending messages.
+		 * This prevents concurrent modifications that could lead to corrupted SSE events.
+		 */
+		private final ReentrantLock sseBuilderLock = new ReentrantLock();
+
 		/**
 		 * Creates a new session transport with the specified ID and SSE builder.
 		 * @param sessionId The unique identifier for this session
@@ -358,6 +365,7 @@ private class WebMvcMcpSessionTransport implements McpServerTransport {
 		@Override
 		public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message) {
 			return Mono.fromRunnable(() -> {
+				sseBuilderLock.lock();
 				try {
 					String jsonText = objectMapper.writeValueAsString(message);
 					sseBuilder.id(sessionId).event(MESSAGE_EVENT_TYPE).data(jsonText);
@@ -367,6 +375,9 @@ public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message) {
 					logger.error("Failed to send message to session {}: {}", sessionId, e.getMessage());
 					sseBuilder.error(e);
 				}
+				finally {
+					sseBuilderLock.unlock();
+				}
 			});
 		}
 
@@ -390,13 +401,17 @@ public <T> T unmarshalFrom(Object data, TypeReference<T> typeRef) {
 		public Mono<Void> closeGracefully() {
 			return Mono.fromRunnable(() -> {
 				logger.debug("Closing session transport: {}", sessionId);
+				sseBuilderLock.lock();
 				try {
 					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 {
+					sseBuilderLock.unlock();
+				}
 			});
 		}
 
@@ -405,13 +420,17 @@ public Mono<Void> closeGracefully() {
 		 */
 		@Override
 		public void close() {
+			sseBuilderLock.lock();
 			try {
 				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 {
+				sseBuilderLock.unlock();
+			}
 		}
 
 	}