feat: add configurable keep-alive support to streamable server transports

- Add KeepAliveScheduler integration to WebFlux and WebMVC transport providers
- Support configurable keep-alive intervals through builder pattern
- Automatically shutdown keep-alive scheduler during graceful shutdown

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

Author: tzolov

mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/server/transport/WebFluxStreamableServerTransportProvider.java

@@ -12,6 +12,8 @@
 import io.modelcontextprotocol.spec.McpStreamableServerTransportProvider;
 import io.modelcontextprotocol.server.McpTransportContext;
 import io.modelcontextprotocol.util.Assert;
+import io.modelcontextprotocol.util.KeepAliveScheduler;
+
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.http.HttpStatus;
@@ -28,6 +30,7 @@
 import reactor.core.publisher.Mono;
 
 import java.io.IOException;
+import java.time.Duration;
 import java.util.List;
 import java.util.concurrent.ConcurrentHashMap;
 
@@ -58,8 +61,11 @@ public class WebFluxStreamableServerTransportProvider implements McpStreamableSe
 
 	private volatile boolean isClosing = false;
 
+	private KeepAliveScheduler keepAliveScheduler;
+
 	private WebFluxStreamableServerTransportProvider(ObjectMapper objectMapper, String mcpEndpoint,
-			McpTransportContextExtractor<ServerRequest> contextExtractor, boolean disallowDelete) {
+			McpTransportContextExtractor<ServerRequest> contextExtractor, boolean disallowDelete,
+			Duration keepAliveInterval) {
 		Assert.notNull(objectMapper, "ObjectMapper must not be null");
 		Assert.notNull(mcpEndpoint, "Message endpoint must not be null");
 		Assert.notNull(contextExtractor, "Context extractor must not be null");
@@ -73,6 +79,20 @@ private WebFluxStreamableServerTransportProvider(ObjectMapper objectMapper, Stri
 			.POST(this.mcpEndpoint, this::handlePost)
 			.DELETE(this.mcpEndpoint, this::handleDelete)
 			.build();
+
+		if (keepAliveInterval != null) {
+			this.keepAliveScheduler = KeepAliveScheduler
+				.builder(() -> (isClosing) ? Flux.empty() : Flux.fromIterable(this.sessions.values()))
+				.initialDelay(keepAliveInterval)
+				.interval(keepAliveInterval)
+				.build();
+
+			this.keepAliveScheduler.start();
+		}
+		else {
+			logger.warn("Keep-alive interval is not set or invalid. No keep-alive will be scheduled.");
+		}
+
 	}
 
 	@Override
@@ -105,6 +125,11 @@ public Mono<Void> closeGracefully() {
 				.doFirst(() -> logger.debug("Initiating graceful shutdown with {} active sessions", sessions.size()))
 				.flatMap(McpStreamableServerSession::closeGracefully)
 				.then();
+		}).then().doOnSuccess(v -> {
+			sessions.clear();
+			if (this.keepAliveScheduler != null) {
+				this.keepAliveScheduler.shutdown();
+			}
 		});
 	}
 
@@ -368,6 +393,8 @@ public static class Builder {
 
 		private boolean disallowDelete;
 
+		private Duration keepAliveInterval;
+
 		private Builder() {
 			// used by a static method
 		}
@@ -424,6 +451,17 @@ public Builder disallowDelete(boolean disallowDelete) {
 			return this;
 		}
 
+		/**
+		 * Sets the keep-alive interval for the server transport.
+		 * @param keepAliveInterval The interval for sending keep-alive messages. If null,
+		 * no keep-alive will be scheduled.
+		 * @return this builder instance
+		 */
+		public Builder keepAliveInterval(Duration keepAliveInterval) {
+			this.keepAliveInterval = keepAliveInterval;
+			return this;
+		}
+
 		/**
 		 * Builds a new instance of {@link WebFluxStreamableServerTransportProvider} with
 		 * the configured settings.
@@ -435,7 +473,7 @@ public WebFluxStreamableServerTransportProvider build() {
 			Assert.notNull(mcpEndpoint, "Message endpoint must be set");
 
 			return new WebFluxStreamableServerTransportProvider(objectMapper, mcpEndpoint, contextExtractor,
-					disallowDelete);
+					disallowDelete, keepAliveInterval);
 		}
 
 	}

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

@@ -33,6 +33,8 @@
 import io.modelcontextprotocol.spec.McpStreamableServerTransport;
 import io.modelcontextprotocol.spec.McpStreamableServerTransportProvider;
 import io.modelcontextprotocol.util.Assert;
+import io.modelcontextprotocol.util.KeepAliveScheduler;
+import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
 /**
@@ -101,6 +103,8 @@ public class WebMvcStreamableServerTransportProvider implements McpStreamableSer
 	 */
 	private volatile boolean isClosing = false;
 
+	private KeepAliveScheduler keepAliveScheduler;
+
 	/**
 	 * Constructs a new WebMvcStreamableServerTransportProvider instance.
 	 * @param objectMapper The ObjectMapper to use for JSON serialization/deserialization
@@ -113,7 +117,8 @@ public class WebMvcStreamableServerTransportProvider implements McpStreamableSer
 	 * @throws IllegalArgumentException if any parameter is null
 	 */
 	private WebMvcStreamableServerTransportProvider(ObjectMapper objectMapper, String mcpEndpoint,
-			boolean disallowDelete, McpTransportContextExtractor<ServerRequest> contextExtractor) {
+			boolean disallowDelete, McpTransportContextExtractor<ServerRequest> contextExtractor,
+			Duration keepAliveInterval) {
 		Assert.notNull(objectMapper, "ObjectMapper must not be null");
 		Assert.notNull(mcpEndpoint, "MCP endpoint must not be null");
 		Assert.notNull(contextExtractor, "McpTransportContextExtractor must not be null");
@@ -127,6 +132,19 @@ private WebMvcStreamableServerTransportProvider(ObjectMapper objectMapper, Strin
 			.POST(this.mcpEndpoint, this::handlePost)
 			.DELETE(this.mcpEndpoint, this::handleDelete)
 			.build();
+
+		if (keepAliveInterval != null) {
+			this.keepAliveScheduler = KeepAliveScheduler
+				.builder(() -> (isClosing) ? Flux.empty() : Flux.fromIterable(this.sessions.values()))
+				.initialDelay(keepAliveInterval)
+				.interval(keepAliveInterval)
+				.build();
+
+			this.keepAliveScheduler.start();
+		}
+		else {
+			logger.warn("Keep-alive interval is not set or invalid. No keep-alive will be scheduled.");
+		}
 	}
 
 	@Override
@@ -184,6 +202,12 @@ public Mono<Void> closeGracefully() {
 
 			this.sessions.clear();
 			logger.debug("Graceful shutdown completed");
+		}).then().doOnSuccess(v -> {
+			logger.debug("Graceful shutdown completed");
+			sessions.clear();
+			if (this.keepAliveScheduler != null) {
+				this.keepAliveScheduler.shutdown();
+			}
 		});
 	}
 
@@ -584,6 +608,8 @@ public static class Builder {
 
 		private McpTransportContextExtractor<ServerRequest> contextExtractor = (serverRequest, context) -> context;
 
+		private Duration keepAliveInterval;
+
 		/**
 		 * Sets the ObjectMapper to use for JSON serialization/deserialization of MCP
 		 * messages.
@@ -635,6 +661,18 @@ public Builder contextExtractor(McpTransportContextExtractor<ServerRequest> cont
 			return this;
 		}
 
+		/**
+		 * Sets the keep-alive interval for the transport. If set, a keep-alive scheduler
+		 * will be created to periodically check and send keep-alive messages to clients.
+		 * @param keepAliveInterval The interval duration for keep-alive messages, or null
+		 * to disable keep-alive
+		 * @return this builder instance
+		 */
+		public Builder keepAliveInterval(Duration keepAliveInterval) {
+			this.keepAliveInterval = keepAliveInterval;
+			return this;
+		}
+
 		/**
 		 * Builds a new instance of {@link WebMvcStreamableServerTransportProvider} with
 		 * the configured settings.
@@ -646,7 +684,7 @@ public WebMvcStreamableServerTransportProvider build() {
 			Assert.notNull(this.mcpEndpoint, "MCP endpoint must be set");
 
 			return new WebMvcStreamableServerTransportProvider(this.objectMapper, this.mcpEndpoint, this.disallowDelete,
-					this.contextExtractor);
+					this.contextExtractor, this.keepAliveInterval);
 		}
 
 	}

mcp/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerSession.java

@@ -1,26 +1,27 @@
 package io.modelcontextprotocol.spec;
 
+import java.time.Duration;
+import java.util.Map;
+import java.util.UUID;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Supplier;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
 import com.fasterxml.jackson.core.type.TypeReference;
+
 import io.modelcontextprotocol.server.McpAsyncServerExchange;
 import io.modelcontextprotocol.server.McpNotificationHandler;
 import io.modelcontextprotocol.server.McpRequestHandler;
 import io.modelcontextprotocol.server.McpTransportContext;
 import io.modelcontextprotocol.util.Assert;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.core.publisher.MonoSink;
 
-import java.time.Duration;
-import java.util.Map;
-import java.util.UUID;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.atomic.AtomicInteger;
-import java.util.concurrent.atomic.AtomicLong;
-import java.util.concurrent.atomic.AtomicReference;
-import java.util.function.Supplier;
-
 /**
  * Representation of a Streamable HTTP server session that keeps track of mapping
  * server-initiated requests to the client and mapping arriving responses. It also allows