KeepAliveScheduler only pings sessions with an open stream

Signed-off-by: Daniel Garnier-Moiroux <git@garnier.wf>

mcp-core/src/main/java/io/modelcontextprotocol/server/transport/HttpServletStreamableServerTransportProvider.java

@@ -20,6 +20,7 @@ import io.modelcontextprotocol.server.McpTransportContextExtractor;
 import io.modelcontextprotocol.spec.HttpHeaders;
 import io.modelcontextprotocol.spec.McpError;
 import io.modelcontextprotocol.spec.McpSchema;
+import io.modelcontextprotocol.spec.McpSession;
 import io.modelcontextprotocol.spec.McpStreamableServerSession;
 import io.modelcontextprotocol.spec.McpStreamableServerTransport;
 import io.modelcontextprotocol.spec.McpStreamableServerTransportProvider;
@@ -165,8 +166,7 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet
 
 		if (keepAliveInterval != null) {
 
-			this.keepAliveScheduler = KeepAliveScheduler
+			this.keepAliveScheduler = KeepAliveScheduler.builder(this::sessionsToPing)
-				.builder(() -> (isClosing) ? Flux.empty() : Flux.fromIterable(sessions.values()))
 				.initialDelay(keepAliveInterval)
 				.interval(keepAliveInterval)
 				.build();
@@ -176,6 +176,22 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet
 
 	}
 
+	/**
+	 * Returns the sessions a keep-alive ping can be sent to, that is the sessions having
+	 * a listening stream. A session without one, e.g. a client which only ever issues
+	 * POST requests, has nothing to write a ping to: pinging it would fail on every
+	 * interval without ever telling us anything about the client being alive.
+	 * @return the sessions to ping
+	 */
+	private Flux<McpSession> sessionsToPing() {
+		if (this.isClosing) {
+			return Flux.empty();
+		}
+		return Flux.fromIterable(this.sessions.values())
+			.filter(McpStreamableServerSession::hasListeningStream)
+			.cast(McpSession.class);
+	}
+
 	@Override
 	public void setSessionFactory(McpStreamableServerSession.Factory sessionFactory) {
 		this.sessionFactory = sessionFactory;

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

@@ -178,6 +178,17 @@ public class McpStreamableServerSession implements McpLoggableSession {
 		}));
 	}
 
+	/**
+	 * Whether the session currently has a listening stream, that is a stream the server
+	 * can send its own requests and notifications to. Sessions have none until the client
+	 * issues the GET request establishing one, and clients are not required to ever issue
+	 * it.
+	 * @return {@code true} if the session has a listening stream
+	 */
+	public boolean hasListeningStream() {
+		return this.listeningStreamRef.get() instanceof McpStreamableServerSessionStream;
+	}
+
 	/**
 	 * Create a listening stream (the generic HTTP GET request, with or without a
 	 * Last-Event-ID header). A session addresses a single listening stream at a time, so

mcp-test/src/test/java/io/modelcontextprotocol/server/HttpServletStreamableIntegrationTests.java

@@ -20,6 +20,9 @@ import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
 import java.util.stream.Stream;
 
+import ch.qos.logback.classic.Level;
+import ch.qos.logback.classic.spi.ILoggingEvent;
+import ch.qos.logback.core.read.ListAppender;
 import io.modelcontextprotocol.AbstractMcpClientServerIntegrationTests;
 import io.modelcontextprotocol.client.McpClient;
 import io.modelcontextprotocol.client.transport.HttpClientStreamableHttpTransport;
@@ -30,6 +33,7 @@ import io.modelcontextprotocol.server.transport.HttpServletStreamableServerTrans
 import io.modelcontextprotocol.server.transport.TomcatTestUtil;
 import io.modelcontextprotocol.spec.HttpHeaders;
 import io.modelcontextprotocol.spec.McpSchema;
+import io.modelcontextprotocol.util.KeepAliveScheduler;
 import jakarta.servlet.http.HttpServletRequest;
 import jakarta.servlet.http.HttpServletResponse;
 import org.apache.catalina.LifecycleException;
@@ -42,6 +46,7 @@ import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.Timeout;
 import org.junit.jupiter.params.provider.Arguments;
+import org.slf4j.LoggerFactory;
 import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
 
@@ -107,7 +112,7 @@ class HttpServletStreamableIntegrationTests extends AbstractMcpClientServerInteg
 		mcpServerTransportProvider = HttpServletStreamableServerTransportProvider.builder()
 			.contextExtractor(TEST_CONTEXT_EXTRACTOR)
 			.mcpEndpoint(MESSAGE_ENDPOINT)
-			.keepAliveInterval(Duration.ofSeconds(1))
+			.keepAliveInterval(Duration.ofMillis(200))
 			.maxRequestSize(MAX_REQUEST_SIZE)
 			.build();
 		MCP_SERVLET.setDelegate(mcpServerTransportProvider);
@@ -267,6 +272,46 @@ class HttpServletStreamableIntegrationTests extends AbstractMcpClientServerInteg
 		});
 	}
 
+	@Test
+	void keepAliveSkipsSessionsWithoutListeningStream() throws Exception {
+		prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0").build();
+		var httpClient = HttpClient.newHttpClient();
+		var keepAliveLogger = (ch.qos.logback.classic.Logger) LoggerFactory.getLogger(KeepAliveScheduler.class);
+		ListAppender<ILoggingEvent> logAppender = new ListAppender<>();
+		logAppender.start();
+		keepAliveLogger.addAppender(logAppender);
+
+		try {
+			// A client is free to never issue the GET request establishing a listening
+			// stream. Such a session has nothing to write a ping to, so it must not be
+			// pinged on every keep-alive interval.
+			initializeSession(httpClient);
+
+			// The keep-alive interval is 200ms, so this spans several intervals
+			Thread.sleep(1_000);
+
+			assertThat(logAppender.list).noneMatch(event -> event.getLevel() == Level.WARN);
+		}
+		finally {
+			keepAliveLogger.detachAppender(logAppender);
+			logAppender.stop();
+		}
+	}
+
+	@Test
+	void keepAlivePingsSessionsWithListeningStream() throws Exception {
+		prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0").build();
+		var httpClient = HttpClient.newHttpClient();
+
+		var sessionId = initializeSession(httpClient);
+		var stream = openListeningStream(httpClient, sessionId, null);
+		awaitStreamOpen(stream);
+
+		// Sessions with a listening stream are still pinged
+		await().atMost(Duration.ofSeconds(1))
+			.untilAsserted(() -> assertThat(stream.events()).anyMatch(line -> line.contains("\"method\":\"ping\"")));
+	}
+
 	private String initializeSession(HttpClient httpClient) throws Exception {
 		var initialize = HttpRequest.newBuilder()
 			.uri(URI.create("http://localhost:" + PORT + MESSAGE_ENDPOINT))