Fix integration tests for roots/list

There is a race condition where the GET /mcp SSE stream in Streamable
HTTP is established _after_ the server tries to send notifications. This
surfaces errors from the MCP client sending notifications and breaks
some integration tests.

We add a utility showing to ensure the stream is established before
sending server notifications.

Fixes #1114

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

Author: Kehrlann

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

@@ -82,6 +82,13 @@ public abstract class AbstractMcpClientServerIntegrationTests {
 
 	abstract protected McpServer.SyncSpecification<?> prepareSyncServerBuilder();
 
+	// There is, for Streamable HTTP, a race condition between establishing the SSE stream
+	// and the server sending notifications. This breaks some `roots/list` tests (and
+	// could in theory break sampling and elicitation tests). This utility method allows
+	// delaying the test until the stream is opened.
+	protected void awaitClientStreamEstablished() {
+	}
+
 	@ParameterizedTest(name = "{0} : {displayName} ")
 	@MethodSource("clientsForTesting")
 	void simple(String clientType) {
@@ -1082,6 +1089,7 @@ void testRootsSuccess(String clientType) {
 
 			InitializeResult initResult = mcpClient.initialize();
 			assertThat(initResult).isNotNull();
+			awaitClientStreamEstablished();
 
 			assertThat(rootsRef.get()).isNull();
 
@@ -1168,7 +1176,7 @@ void testRootsNotificationWithEmptyRootsList(String clientType) {
 
 			InitializeResult initResult = mcpClient.initialize();
 			assertThat(initResult).isNotNull();
-
+			awaitClientStreamEstablished();
 			mcpClient.rootsListChangedNotification();
 
 			await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
@@ -1201,7 +1209,7 @@ void testRootsWithMultipleHandlers(String clientType) {
 			.build()) {
 
 			assertThat(mcpClient.initialize()).isNotNull();
-
+			awaitClientStreamEstablished();
 			mcpClient.rootsListChangedNotification();
 
 			await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
@@ -1234,7 +1242,7 @@ void testRootsServerCloseWithActiveSubscription(String clientType) {
 
 			InitializeResult initResult = mcpClient.initialize();
 			assertThat(initResult).isNotNull();
-
+			awaitClientStreamEstablished();
 			mcpClient.rootsListChangedNotification();
 
 			await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {

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

@@ -30,7 +30,6 @@
 import org.apache.catalina.LifecycleException;
 import org.apache.catalina.LifecycleState;
 import org.apache.catalina.startup.Tomcat;
-import org.awaitility.Awaitility;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeAll;
@@ -42,6 +41,7 @@
 import reactor.test.StepVerifier;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.awaitility.Awaitility.await;
 
 @Timeout(15)
 class HttpServletStreamableIntegrationTests extends AbstractMcpClientServerIntegrationTests {
@@ -57,6 +57,16 @@ class HttpServletStreamableIntegrationTests extends AbstractMcpClientServerInteg
 
 	private HttpServletStreamableServerTransportProvider mcpServerTransportProvider;
 
+	@Override
+	protected void awaitClientStreamEstablished() {
+		var timeout = Duration.ofSeconds(1);
+		await().atMost(timeout).untilAsserted(() -> {
+			assertThat(MCP_SERVLET.isStreamEstablished())
+				.withFailMessage("[Failed to observe MCP Client connection within %s]", timeout)
+				.isTrue();
+		});
+	}
+
 	static Stream<Arguments> clientsForTesting() {
 		return Stream.of(Arguments.of("httpclient"));
 	}
@@ -151,7 +161,7 @@ void testMissingHandlerReturnsMethodNotFoundError() {
 				.verifyComplete();
 
 			// Wait until we've received the response
-			Awaitility.await().atMost(Duration.ofSeconds(1)).until(() -> response.get() != null);
+			await().atMost(Duration.ofSeconds(1)).until(() -> response.get() != null);
 
 			assertThat(response.get().error().code()).isEqualTo(McpSchema.ErrorCodes.METHOD_NOT_FOUND);
 			assertThat(response.get().error().message()).isEqualTo("Method not found: foo/bar");

mcp-test/src/test/java/io/modelcontextprotocol/server/transport/TomcatTestUtil.java

@@ -7,14 +7,17 @@
 import java.io.IOException;
 import java.net.InetSocketAddress;
 import java.net.ServerSocket;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import jakarta.servlet.Filter;
 import jakarta.servlet.Servlet;
 import jakarta.servlet.ServletConfig;
 import jakarta.servlet.ServletException;
 import jakarta.servlet.ServletRequest;
 import jakarta.servlet.ServletResponse;
+import jakarta.servlet.http.HttpServletRequest;
 import org.apache.catalina.Context;
+import org.apache.catalina.Wrapper;
 import org.apache.catalina.startup.Tomcat;
 import org.apache.tomcat.util.descriptor.web.FilterDef;
 import org.apache.tomcat.util.descriptor.web.FilterMap;
@@ -41,7 +44,7 @@ public static Tomcat createTomcatServer(String contextPath, int port, Servlet se
 		Context context = tomcat.addContext(contextPath, baseDir);
 
 		// Add transport servlet to Tomcat
-		org.apache.catalina.Wrapper wrapper = context.createWrapper();
+		Wrapper wrapper = context.createWrapper();
 		wrapper.setName("mcpServlet");
 		wrapper.setServlet(servlet);
 		wrapper.setLoadOnStartup(1);
@@ -78,12 +81,18 @@ public static class DelegatingServlet implements Servlet {
 
 		private volatile Servlet delegate;
 
+		// Crude way of tracking whether a GET SSE stream has been
+		// established, to ensure a Streamable HTTP MCP Client is
+		// connected.
+		private final AtomicBoolean sseStreamEstablished = new AtomicBoolean(false);
+
 		/**
 		 * Sets the servlet handling subsequent requests. The delegate is not
 		 * {@link Servlet#init(ServletConfig) initialized}, since the MCP servlet
 		 * transports do not rely on their {@link ServletConfig}.
 		 */
 		public void setDelegate(Servlet delegate) {
+			this.sseStreamEstablished.set(false);
 			this.delegate = delegate;
 		}
 
@@ -104,6 +113,9 @@ public void service(ServletRequest request, ServletResponse response) throws Ser
 				throw new IllegalStateException("No delegate servlet has been set");
 			}
 			current.service(request, response);
+			if (request instanceof HttpServletRequest req && req.getMethod().equals("GET")) {
+				sseStreamEstablished.set(true);
+			}
 		}
 
 		@Override
@@ -116,6 +128,10 @@ public void destroy() {
 			this.delegate = null;
 		}
 
+		public boolean isStreamEstablished() {
+			return this.sseStreamEstablished.get();
+		}
+
 	}
 
 	/**