optimize and add test for listening sse closed

Author: taobaorun

mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/WebMvcStreamableIntegrationTests.java

@@ -81,17 +81,24 @@ public void before() {
 			throw new RuntimeException("Failed to start Tomcat", e);
 		}
 
-		clientBuilders
-			.put("httpclient",
-					McpClient.sync(HttpClientStreamableHttpTransport.builder("http://localhost:" + PORT)
-						.endpoint(MESSAGE_ENDPOINT)
-						.build()).initializationTimeout(Duration.ofHours(10)).requestTimeout(Duration.ofHours(10)));
-
-		clientBuilders.put("webflux",
-				McpClient.sync(WebClientStreamableHttpTransport
-					.builder(WebClient.builder().baseUrl("http://localhost:" + PORT))
-					.endpoint(MESSAGE_ENDPOINT)
-					.build()));
+		var httpClientTransport = HttpClientStreamableHttpTransport.builder("http://localhost:" + PORT)
+			.endpoint(MESSAGE_ENDPOINT)
+			.openConnectionOnStartup(true)
+			.build();
+
+		clientTransportBuilders.put("httpclient", httpClientTransport);
+		clientBuilders.put("httpclient",
+				McpClient.sync(httpClientTransport)
+					.initializationTimeout(Duration.ofHours(10))
+					.requestTimeout(Duration.ofHours(10)));
+		var webClientTransport = WebClientStreamableHttpTransport
+			.builder(WebClient.builder().baseUrl("http://localhost:" + PORT))
+			.endpoint(MESSAGE_ENDPOINT)
+			.openConnectionOnStartup(true)
+			.build();
+		clientTransportBuilders.put("webflux", webClientTransport);
+
+		clientBuilders.put("webflux", McpClient.sync(webClientTransport));
 
 		// Get the transport from Spring context
 		this.mcpServerTransportProvider = tomcatServer.appContext()

mcp-spring/mcp-spring-webmvc/src/test/resources/logback.xml

@@ -18,7 +18,7 @@
     <logger name="io.modelcontextprotocol.server.transport" level="INFO"/>
 
     <!-- Spec package -->
-    <logger name="io.modelcontextprotocol.spec" level="INFO"/>
+    <logger name="io.modelcontextprotocol.spec" level="DEBUG"/>
 
     <!-- Root logger -->
     <root level="INFO">

mcp-test/src/main/java/io/modelcontextprotocol/AbstractMcpClientServerIntegrationTests.java

@@ -12,6 +12,9 @@
 import static org.awaitility.Awaitility.await;
 import static org.mockito.Mockito.mock;
 
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.PrintStream;
 import java.net.URI;
 import java.net.http.HttpClient;
 import java.net.http.HttpRequest;
@@ -29,6 +32,7 @@
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
+import io.modelcontextprotocol.spec.McpClientTransport;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
 
@@ -66,6 +70,8 @@ public abstract class AbstractMcpClientServerIntegrationTests {
 
 	protected ConcurrentHashMap<String, McpClient.SyncSpec> clientBuilders = new ConcurrentHashMap<>();
 
+	protected ConcurrentHashMap<String, McpClientTransport> clientTransportBuilders = new ConcurrentHashMap<>();
+
 	abstract protected void prepareClients(int port, String mcpEndpoint);
 
 	abstract protected McpServer.AsyncSpecification<?> prepareAsyncServerBuilder();
@@ -836,7 +842,7 @@ void testThrowingToolCallIsCaughtBeforeTimeout(String clientType) {
 
 	@ParameterizedTest(name = "{0} : {displayName} ")
 	@ValueSource(strings = { "httpclient", "webflux" })
-	void testToolCallSuccessWithTranportContextExtraction(String clientType) {
+	void testToolCallSuccessWithTransportContextExtraction(String clientType) {
 
 		var clientBuilder = clientBuilders.get(clientType);
 
@@ -999,6 +1005,32 @@ void testInitialize(String clientType) {
 		mcpServer.close();
 	}
 
+	@ParameterizedTest(name = "{0} : {displayName} ")
+	@ValueSource(strings = { "httpclient", "webflux" })
+	void testListeningStreamWillClosedWhenNew(String clientType) throws IOException {
+		var clientTransport = clientTransportBuilders.get(clientType);
+		if (clientTransport == null) {
+			return;
+		}
+		PrintStream originalOut = System.out;
+		ByteArrayOutputStream capturedOutput = new ByteArrayOutputStream();
+		System.setOut(new PrintStream(capturedOutput));
+
+		var clientBuilder = clientBuilders.get(clientType);
+		var mcpServer = prepareSyncServerBuilder().build();
+		var mcpClient = clientBuilder.build();
+		InitializeResult initResult = mcpClient.initialize();
+		assertThat(initResult).isNotNull();
+		clientTransport.connect(message -> Mono.empty()).subscribe();
+		await().atMost(Duration.ofSeconds(1)).untilAsserted(() -> {
+			assertThat(capturedOutput.toString().contains("Listening stream already exists for this session")).isTrue();
+		});
+		System.setOut(originalOut);
+		capturedOutput.close();
+		mcpClient.close();
+		mcpServer.close();
+	}
+
 	// ---------------------------------------
 	// Logging Tests
 	// ---------------------------------------

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

@@ -139,7 +139,7 @@ public Mono<Void> delete() {
 	public Mono<McpStreamableServerSessionStream> listeningStream(McpStreamableServerTransport transport) {
 		McpStreamableServerSessionStream listeningStream = new McpStreamableServerSessionStream(transport);
 		McpLoggableSession oldStream = this.listeningStreamRef.getAndSet(listeningStream);
-		if (oldStream != null) {
+		if (oldStream != null && !(oldStream instanceof MissingMcpTransportSession)) {
 			logger.debug(
 					"Listening stream already exists for this session:{} and will be closed to make way for the new listening SSE stream",
 					this.id);