improve webmvc streamable http server tests

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

Author: tzolov

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

@@ -205,7 +205,7 @@ public RouterFunction<ServerResponse> getRouterFunction() {
 	}
 
 	/**
-	 * Handles GET requests for SSE connections and message replay.
+	 * Setup the listening SSE connections and message replay.
 	 * @param request The incoming server request
 	 * @return A ServerResponse configured for SSE communication, or an error response
 	 */

mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/WebMcpStreamableMcpAsyncServerTests.java renamed to mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/WebMcpStreamableAsyncServerTransportTests.java

@@ -28,7 +28,7 @@
  * @author Christian Tzolov
  */
 @Timeout(15) // Giving extra time beyond the client timeout
-class WebMcpStreamableMcpAsyncServerTests extends AbstractMcpAsyncServerTests {
+class WebMcpStreamableAsyncServerTransportTests extends AbstractMcpAsyncServerTests {
 
 	private static final int PORT = TestUtil.findAvailablePort();
 

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

@@ -0,0 +1,121 @@
+/*
+ * Copyright 2024-2024 the original author or authors.
+ */
+
+package io.modelcontextprotocol.server;
+
+import org.apache.catalina.Context;
+import org.apache.catalina.LifecycleException;
+import org.apache.catalina.startup.Tomcat;
+import org.junit.jupiter.api.Timeout;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.web.context.support.AnnotationConfigWebApplicationContext;
+import org.springframework.web.servlet.DispatcherServlet;
+import org.springframework.web.servlet.config.annotation.EnableWebMvc;
+import org.springframework.web.servlet.function.RouterFunction;
+import org.springframework.web.servlet.function.ServerResponse;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+
+import io.modelcontextprotocol.server.transport.WebMvcStreamableServerTransportProvider;
+import io.modelcontextprotocol.spec.McpStreamableServerTransportProvider;
+import reactor.netty.DisposableServer;
+
+/**
+ * Tests for {@link McpAsyncServer} using {@link WebFluxSseServerTransportProvider}.
+ *
+ * @author Christian Tzolov
+ */
+@Timeout(15) // Giving extra time beyond the client timeout
+class WebMcpStreamableSyncServerTransportTests extends AbstractMcpSyncServerTests {
+
+	private static final int PORT = TestUtil.findAvailablePort();
+
+	private static final String MCP_ENDPOINT = "/mcp";
+
+	private DisposableServer httpServer;
+
+	private AnnotationConfigWebApplicationContext appContext;
+
+	private Tomcat tomcat;
+
+	private McpStreamableServerTransportProvider transportProvider;
+
+	@Configuration
+	@EnableWebMvc
+	static class TestConfig {
+
+		@Bean
+		public WebMvcStreamableServerTransportProvider webMvcSseServerTransportProvider() {
+			return WebMvcStreamableServerTransportProvider.builder()
+				.objectMapper(new ObjectMapper())
+				.mcpEndpoint(MCP_ENDPOINT)
+				.build();
+		}
+
+		@Bean
+		public RouterFunction<ServerResponse> routerFunction(
+				WebMvcStreamableServerTransportProvider transportProvider) {
+			return transportProvider.getRouterFunction();
+		}
+
+	}
+
+	private McpStreamableServerTransportProvider createMcpTransportProvider() {
+		// Set up Tomcat first
+		tomcat = new Tomcat();
+		tomcat.setPort(PORT);
+
+		// Set Tomcat base directory to java.io.tmpdir to avoid permission issues
+		String baseDir = System.getProperty("java.io.tmpdir");
+		tomcat.setBaseDir(baseDir);
+
+		// Use the same directory for document base
+		Context context = tomcat.addContext("", baseDir);
+
+		// Create and configure Spring WebMvc context
+		appContext = new AnnotationConfigWebApplicationContext();
+		appContext.register(TestConfig.class);
+		appContext.setServletContext(context.getServletContext());
+		appContext.refresh();
+
+		// Get the transport from Spring context
+		transportProvider = appContext.getBean(McpStreamableServerTransportProvider.class);
+
+		// Create DispatcherServlet with our Spring context
+		DispatcherServlet dispatcherServlet = new DispatcherServlet(appContext);
+
+		// Add servlet to Tomcat and get the wrapper
+		var wrapper = Tomcat.addServlet(context, "dispatcherServlet", dispatcherServlet);
+		wrapper.setLoadOnStartup(1);
+		context.addServletMappingDecoded("/*", "dispatcherServlet");
+
+		try {
+			tomcat.start();
+			tomcat.getConnector(); // Create and start the connector
+		}
+		catch (LifecycleException e) {
+			throw new RuntimeException("Failed to start Tomcat", e);
+		}
+
+		return transportProvider;
+	}
+
+	@Override
+	protected McpServer.SyncSpecification<?> prepareSyncServerBuilder() {
+		return McpServer.sync(createMcpTransportProvider());
+	}
+
+	@Override
+	protected void onStart() {
+	}
+
+	@Override
+	protected void onClose() {
+		if (httpServer != null) {
+			httpServer.disposeNow();
+		}
+	}
+
+}

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

@@ -18,6 +18,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
@@ -35,6 +36,8 @@
 import io.modelcontextprotocol.spec.McpSchema.ClientCapabilities;
 import io.modelcontextprotocol.spec.McpSchema.CreateMessageRequest;
 import io.modelcontextprotocol.spec.McpSchema.CreateMessageResult;
+import io.modelcontextprotocol.spec.McpSchema.ElicitRequest;
+import io.modelcontextprotocol.spec.McpSchema.ElicitResult;
 import io.modelcontextprotocol.spec.McpSchema.InitializeResult;
 import io.modelcontextprotocol.spec.McpSchema.ModelPreferences;
 import io.modelcontextprotocol.spec.McpSchema.Role;
@@ -489,51 +492,52 @@ void testCreateElicitationWithRequestTimeoutSuccess(String clientType) {
 	@ValueSource(strings = { "httpclient", "webflux" })
 	void testCreateElicitationWithRequestTimeoutFail(String clientType) {
 
+		var latch = new CountDownLatch(1);
+
 		var clientBuilder = clientBuilders.get(clientType);
 
-		Function<McpSchema.ElicitRequest, McpSchema.ElicitResult> elicitationHandler = request -> {
+		Function<ElicitRequest, ElicitResult> elicitationHandler = request -> {
 			assertThat(request.message()).isNotEmpty();
 			assertThat(request.requestedSchema()).isNotNull();
+
 			try {
-				TimeUnit.SECONDS.sleep(2);
+				if (!latch.await(2, TimeUnit.SECONDS)) {
+					throw new RuntimeException("Timeout waiting for elicitation processing");
+				}
 			}
 			catch (InterruptedException e) {
 				throw new RuntimeException(e);
 			}
-			return new McpSchema.ElicitResult(McpSchema.ElicitResult.Action.ACCEPT,
-					Map.of("message", request.message()));
+			return new ElicitResult(ElicitResult.Action.ACCEPT, Map.of("message", request.message()));
 		};
 
 		var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0"))
 			.capabilities(ClientCapabilities.builder().elicitation().build())
 			.elicitation(elicitationHandler)
 			.build();
 
-		CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")),
-				null);
+		CallToolResult callResponse = new CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), null);
+
+		AtomicReference<ElicitResult> resultRef = new AtomicReference<>();
 
 		McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder()
 			.tool(Tool.builder().name("tool1").description("tool1 description").inputSchema(emptyJsonSchema).build())
 			.callHandler((exchange, request) -> {
 
-				var elicitationRequest = McpSchema.ElicitRequest.builder()
+				var elicitationRequest = ElicitRequest.builder()
 					.message("Test message")
 					.requestedSchema(
 							Map.of("type", "object", "properties", Map.of("message", Map.of("type", "string"))))
 					.build();
 
-				StepVerifier.create(exchange.createElicitation(elicitationRequest)).consumeNextWith(result -> {
-					assertThat(result).isNotNull();
-					assertThat(result.action()).isEqualTo(McpSchema.ElicitResult.Action.ACCEPT);
-					assertThat(result.content().get("message")).isEqualTo("Test message");
-				}).verifyComplete();
-
-				return Mono.just(callResponse);
+				return exchange.createElicitation(elicitationRequest)
+					.doOnNext(resultRef::set)
+					.then(Mono.just(callResponse));
 			})
 			.build();
 
 		var mcpServer = prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0")
-			.requestTimeout(Duration.ofSeconds(1))
+			.requestTimeout(Duration.ofSeconds(1)) // 1 second.
 			.tools(tool)
 			.build();
 
@@ -542,7 +546,10 @@ void testCreateElicitationWithRequestTimeoutFail(String clientType) {
 
 		assertThatExceptionOfType(McpError.class).isThrownBy(() -> {
 			mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of()));
-		}).withMessageContaining("Timeout");
+		}).withMessageContaining("within 1000ms");
+
+		ElicitResult elicitResult = resultRef.get();
+		assertThat(elicitResult).isNull();
 
 		mcpClient.closeGracefully();
 		mcpServer.closeGracefully().block();