From e4d9f7780f40b3310dfd98e4e491a812ee8f13e8 Mon Sep 17 00:00:00 2001
From: zhangzhenhua <zhangzhenhua@qiyi.com>
Date: Wed, 2 Apr 2025 13:54:14 +0800
Subject: [PATCH] Compatible with spring.webflux.base-path

---
 .../WebFluxSseServerTransportProvider.java    | 48 +++++++++++++++++--
 1 file changed, 43 insertions(+), 5 deletions(-)

diff --git a/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/server/transport/WebFluxSseServerTransportProvider.java b/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/server/transport/WebFluxSseServerTransportProvider.java
index 85a39a82..de577c17 100644
--- a/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/server/transport/WebFluxSseServerTransportProvider.java
+++ b/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/server/transport/WebFluxSseServerTransportProvider.java
@@ -82,8 +82,15 @@ public class WebFluxSseServerTransportProvider implements McpServerTransportProv
 	 */
 	public static final String DEFAULT_SSE_ENDPOINT = "/sse";
 
+	public static final String DEFAULT_BASE_PATH = "";
+
 	private final ObjectMapper objectMapper;
 
+	/**
+	 * the project base path
+	 */
+	private final String basePath;
+
 	private final String messageEndpoint;
 
 	private final String sseEndpoint;
@@ -102,6 +109,7 @@ public class WebFluxSseServerTransportProvider implements McpServerTransportProv
 	 */
 	private volatile boolean isClosing = false;
 
+
 	/**
 	 * Constructs a new WebFlux SSE server transport provider instance.
 	 * @param objectMapper The ObjectMapper to use for JSON serialization/deserialization
@@ -112,17 +120,33 @@ public class WebFluxSseServerTransportProvider implements McpServerTransportProv
 	 * @throws IllegalArgumentException if either parameter is null
 	 */
 	public WebFluxSseServerTransportProvider(ObjectMapper objectMapper, String messageEndpoint, String sseEndpoint) {
+		this(objectMapper, DEFAULT_BASE_PATH, messageEndpoint, sseEndpoint);
+	}
+
+
+	/**
+	 * Constructs a new WebFlux SSE server transport provider instance.
+	 * @param objectMapper The ObjectMapper to use for JSON serialization/deserialization
+	 * of MCP messages. Must not be null.
+	 * @param basePath webflux base path
+	 * @param messageEndpoint The endpoint URI where clients should send their JSON-RPC
+	 * messages. This endpoint will be communicated to clients during SSE connection
+	 * setup. Must not be null.
+	 * @throws IllegalArgumentException if either parameter is null
+	 */
+	public WebFluxSseServerTransportProvider(ObjectMapper objectMapper, String basePath, String messageEndpoint, String sseEndpoint) {
 		Assert.notNull(objectMapper, "ObjectMapper must not be null");
 		Assert.notNull(messageEndpoint, "Message endpoint must not be null");
 		Assert.notNull(sseEndpoint, "SSE endpoint must not be null");
 
 		this.objectMapper = objectMapper;
+		this.basePath = basePath;
 		this.messageEndpoint = messageEndpoint;
 		this.sseEndpoint = sseEndpoint;
 		this.routerFunction = RouterFunctions.route()
-			.GET(this.sseEndpoint, this::handleSseConnection)
-			.POST(this.messageEndpoint, this::handleMessage)
-			.build();
+				.GET(this.sseEndpoint, this::handleSseConnection)
+				.POST(this.messageEndpoint, this::handleMessage)
+				.build();
 	}
 
 	/**
@@ -245,7 +269,7 @@ private Mono<ServerResponse> handleSseConnection(ServerRequest request) {
 				logger.debug("Sending initial endpoint event to session: {}", sessionId);
 				sink.next(ServerSentEvent.builder()
 					.event(ENDPOINT_EVENT_TYPE)
-					.data(messageEndpoint + "?sessionId=" + sessionId)
+					.data(basePath + messageEndpoint + "?sessionId=" + sessionId)
 					.build());
 				sink.onCancel(() -> {
 					logger.debug("Session {} cancelled", sessionId);
@@ -360,6 +384,8 @@ public static class Builder {
 
 		private ObjectMapper objectMapper;
 
+		private String basePath = DEFAULT_BASE_PATH;
+
 		private String messageEndpoint;
 
 		private String sseEndpoint = DEFAULT_SSE_ENDPOINT;
@@ -377,6 +403,18 @@ public Builder objectMapper(ObjectMapper objectMapper) {
 			return this;
 		}
 
+		/**
+		 * Sets the project basePath as endpoint prefix where clients should send their JSON-RPC messages
+		 * @param basePath the project basePath . Must not be null.
+		 * @return this builder instance
+		 * @throws IllegalArgumentException if basePath is null
+		 */
+		public Builder basePath(String basePath) {
+			Assert.notNull(basePath, "basePath must not be null");
+			this.basePath = basePath;
+			return this;
+		}
+
 		/**
 		 * Sets the endpoint URI where clients should send their JSON-RPC messages.
 		 * @param messageEndpoint The message endpoint URI. Must not be null.
@@ -411,7 +449,7 @@ public WebFluxSseServerTransportProvider build() {
 			Assert.notNull(objectMapper, "ObjectMapper must be set");
 			Assert.notNull(messageEndpoint, "Message endpoint must be set");
 
-			return new WebFluxSseServerTransportProvider(objectMapper, messageEndpoint, sseEndpoint);
+			return new WebFluxSseServerTransportProvider(objectMapper, basePath, messageEndpoint, sseEndpoint);
 		}
 
 	}