feat: implement MCP-compliant keep-alive functionality for server transports
- Add KeepAliveScheduler utility class for configurable periodic session pings - Integrate keep-alive support in WebFlux, WebMVC, and HttpServlet SSE transport providers - Add keepAliveInterval configuration option to all transport provider builders - Deprecate existing constructors in favor of builder pattern with enhanced configuration - Update graceful shutdown to properly clean up keep-alive schedulers - Add unit tests for KeepAliveScheduler functionality Implements MCP specification recommendations for connection health detection: - Configurable ping frequency to suit different network environments - Optional keep-alive (disabled by default) to avoid excessive network overhead - Proper resource cleanup to prevent connection leaks https://modelcontextprotocol.io/specification/2025-06-18/basic/utilities/ping#implementation-considerations Resolves: #414, #158 Replaces #353 Signed-off-by: Christian Tzolov <christian.tzolov@broadcom.com>
This commit is contained in:
parent
ae2fc8683b
commit
ee2c891cca
2 changed files with 954 additions and 35 deletions
|
|
@ -1,6 +1,7 @@
|
|||
package io.modelcontextprotocol.server.transport;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.time.Duration;
|
||||
import java.util.concurrent.ConcurrentHashMap;
|
||||
|
||||
import com.fasterxml.jackson.core.type.TypeReference;
|
||||
|
|
@ -11,6 +12,8 @@ import io.modelcontextprotocol.spec.McpServerSession;
|
|||
import io.modelcontextprotocol.spec.McpServerTransport;
|
||||
import io.modelcontextprotocol.spec.McpServerTransportProvider;
|
||||
import io.modelcontextprotocol.util.Assert;
|
||||
import io.modelcontextprotocol.util.KeepAliveScheduler;
|
||||
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
import reactor.core.Exceptions;
|
||||
|
|
@ -109,6 +112,12 @@ public class WebFluxSseServerTransportProvider implements McpServerTransportProv
|
|||
*/
|
||||
private volatile boolean isClosing = false;
|
||||
|
||||
/**
|
||||
* Keep-alive scheduler for managing session pings. Activated if keepAliveInterval is
|
||||
* set. Disabled by default.
|
||||
*/
|
||||
private KeepAliveScheduler keepAliveScheduler;
|
||||
|
||||
/**
|
||||
* Constructs a new WebFlux SSE server transport provider instance with the default
|
||||
* SSE endpoint.
|
||||
|
|
@ -118,7 +127,10 @@ public class WebFluxSseServerTransportProvider implements McpServerTransportProv
|
|||
* messages. This endpoint will be communicated to clients during SSE connection
|
||||
* setup. Must not be null.
|
||||
* @throws IllegalArgumentException if either parameter is null
|
||||
* @deprecated Use the builder {@link #builder()} instead for better configuration
|
||||
* options.
|
||||
*/
|
||||
@Deprecated
|
||||
public WebFluxSseServerTransportProvider(ObjectMapper objectMapper, String messageEndpoint) {
|
||||
this(objectMapper, messageEndpoint, DEFAULT_SSE_ENDPOINT);
|
||||
}
|
||||
|
|
@ -131,7 +143,10 @@ public class WebFluxSseServerTransportProvider implements McpServerTransportProv
|
|||
* messages. This endpoint will be communicated to clients during SSE connection
|
||||
* setup. Must not be null.
|
||||
* @throws IllegalArgumentException if either parameter is null
|
||||
* @deprecated Use the builder {@link #builder()} instead for better configuration
|
||||
* options.
|
||||
*/
|
||||
@Deprecated
|
||||
public WebFluxSseServerTransportProvider(ObjectMapper objectMapper, String messageEndpoint, String sseEndpoint) {
|
||||
this(objectMapper, DEFAULT_BASE_URL, messageEndpoint, sseEndpoint);
|
||||
}
|
||||
|
|
@ -145,9 +160,32 @@ public class WebFluxSseServerTransportProvider implements McpServerTransportProv
|
|||
* messages. This endpoint will be communicated to clients during SSE connection
|
||||
* setup. Must not be null.
|
||||
* @throws IllegalArgumentException if either parameter is null
|
||||
* @deprecated Use the builder {@link #builder()} instead for better configuration
|
||||
* options.
|
||||
*/
|
||||
@Deprecated
|
||||
public WebFluxSseServerTransportProvider(ObjectMapper objectMapper, String baseUrl, String messageEndpoint,
|
||||
String sseEndpoint) {
|
||||
this(objectMapper, baseUrl, messageEndpoint, sseEndpoint, null);
|
||||
}
|
||||
|
||||
/**
|
||||
* 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 baseUrl webflux message 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.
|
||||
* @param sseEndpoint The SSE endpoint path. Must not be null.
|
||||
* @param keepAliveInterval The interval for sending keep-alive pings to clients.
|
||||
* @throws IllegalArgumentException if either parameter is null
|
||||
* @deprecated Use the builder {@link #builder()} instead for better configuration
|
||||
* options.
|
||||
*/
|
||||
@Deprecated
|
||||
public WebFluxSseServerTransportProvider(ObjectMapper objectMapper, String baseUrl, String messageEndpoint,
|
||||
String sseEndpoint, Duration keepAliveInterval) {
|
||||
Assert.notNull(objectMapper, "ObjectMapper must not be null");
|
||||
Assert.notNull(baseUrl, "Message base path must not be null");
|
||||
Assert.notNull(messageEndpoint, "Message endpoint must not be null");
|
||||
|
|
@ -161,6 +199,17 @@ public class WebFluxSseServerTransportProvider implements McpServerTransportProv
|
|||
.GET(this.sseEndpoint, this::handleSseConnection)
|
||||
.POST(this.messageEndpoint, this::handleMessage)
|
||||
.build();
|
||||
|
||||
if (keepAliveInterval != null) {
|
||||
|
||||
this.keepAliveScheduler = KeepAliveScheduler
|
||||
.builder(() -> (isClosing) ? Flux.empty() : Flux.fromIterable(sessions.values()))
|
||||
.initialDelay(keepAliveInterval)
|
||||
.interval(keepAliveInterval)
|
||||
.build();
|
||||
|
||||
this.keepAliveScheduler.start();
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
|
|
@ -209,15 +258,6 @@ public class WebFluxSseServerTransportProvider implements McpServerTransportProv
|
|||
/**
|
||||
* Initiates a graceful shutdown of all the sessions. This method ensures all active
|
||||
* sessions are properly closed and cleaned up.
|
||||
*
|
||||
* <p>
|
||||
* The shutdown process:
|
||||
* <ul>
|
||||
* <li>Marks the transport as closing to prevent new connections</li>
|
||||
* <li>Closes each active session</li>
|
||||
* <li>Removes closed sessions from the sessions map</li>
|
||||
* <li>Times out after 5 seconds if shutdown takes too long</li>
|
||||
* </ul>
|
||||
* @return A Mono that completes when all sessions have been closed
|
||||
*/
|
||||
@Override
|
||||
|
|
@ -225,7 +265,14 @@ public class WebFluxSseServerTransportProvider implements McpServerTransportProv
|
|||
return Flux.fromIterable(sessions.values())
|
||||
.doFirst(() -> logger.debug("Initiating graceful shutdown with {} active sessions", sessions.size()))
|
||||
.flatMap(McpServerSession::closeGracefully)
|
||||
.then();
|
||||
.then()
|
||||
.doOnSuccess(v -> {
|
||||
logger.debug("Graceful shutdown completed");
|
||||
sessions.clear();
|
||||
if (this.keepAliveScheduler != null) {
|
||||
this.keepAliveScheduler.shutdown();
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -396,6 +443,8 @@ public class WebFluxSseServerTransportProvider implements McpServerTransportProv
|
|||
|
||||
private String sseEndpoint = DEFAULT_SSE_ENDPOINT;
|
||||
|
||||
private Duration keepAliveInterval;
|
||||
|
||||
/**
|
||||
* Sets the ObjectMapper to use for JSON serialization/deserialization of MCP
|
||||
* messages.
|
||||
|
|
@ -446,6 +495,17 @@ public class WebFluxSseServerTransportProvider implements McpServerTransportProv
|
|||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Sets the interval for sending keep-alive pings to clients.
|
||||
* @param keepAliveInterval The keep-alive interval duration. If null, keep-alive
|
||||
* is disabled.
|
||||
* @return this builder instance
|
||||
*/
|
||||
public Builder keepAliveInterval(Duration keepAliveInterval) {
|
||||
this.keepAliveInterval = keepAliveInterval;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Builds a new instance of {@link WebFluxSseServerTransportProvider} with the
|
||||
* configured settings.
|
||||
|
|
@ -456,7 +516,8 @@ public class WebFluxSseServerTransportProvider implements McpServerTransportProv
|
|||
Assert.notNull(objectMapper, "ObjectMapper must be set");
|
||||
Assert.notNull(messageEndpoint, "Message endpoint must be set");
|
||||
|
||||
return new WebFluxSseServerTransportProvider(objectMapper, baseUrl, messageEndpoint, sseEndpoint);
|
||||
return new WebFluxSseServerTransportProvider(objectMapper, baseUrl, messageEndpoint, sseEndpoint,
|
||||
keepAliveInterval);
|
||||
}
|
||||
|
||||
}
|
||||
|
|
|
|||
|
|
@ -12,6 +12,8 @@ import io.modelcontextprotocol.spec.McpStreamableServerTransport;
|
|||
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.FluxSink;
|
|||
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 @@ public class WebFluxStreamableServerTransportProvider implements McpStreamableSe
|
|||
.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 class WebFluxStreamableServerTransportProvider implements McpStreamableSe
|
|||
.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 class WebFluxStreamableServerTransportProvider implements McpStreamableSe
|
|||
|
||||
private boolean disallowDelete;
|
||||
|
||||
private Duration keepAliveInterval;
|
||||
|
||||
private Builder() {
|
||||
// used by a static method
|
||||
}
|
||||
|
|
@ -424,6 +451,17 @@ public class WebFluxStreamableServerTransportProvider implements McpStreamableSe
|
|||
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 class WebFluxStreamableServerTransportProvider implements McpStreamableSe
|
|||
Assert.notNull(mcpEndpoint, "Message endpoint must be set");
|
||||
|
||||
return new WebFluxStreamableServerTransportProvider(objectMapper, mcpEndpoint, contextExtractor,
|
||||
disallowDelete);
|
||||
disallowDelete, keepAliveInterval);
|
||||
}
|
||||
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue