From edbd78465c37fbce9c89f0c904f0d71efb8aaa1b Mon Sep 17 00:00:00 2001 From: Christian Tzolov Date: Sat, 11 Jan 2025 14:55:52 +0100 Subject: [PATCH] refactor(transport): Replace SseEmitter with queue-based SSE implementation in WebMvc transport Replaces Spring's SseEmitter with a custom BlockingQueue-based SSE implementation in WebMvcSseServerTransport for improved event delivery control and connection management. The new implementation: - Uses BlockingQueue for reliable event queuing and delivery - Adds proper session management with dedicated SSEEvent record - Improves error handling and timeout management - Adds comprehensive integration tests Resolves #57 --- .../mcp-webflux-sse-transport/pom.xml | 9 +- ...s.java => WebFluxSseIntegrationTests.java} | 81 ++- .../transport/WebMvcSseServerTransport.java | 142 +++-- .../mcp/server/WebMvcSseIntegrationTests.java | 524 ++++++++++++++++++ .../src/test/resources/logback.xml | 9 +- .../HttpClientSseClientTransport.java | 4 +- .../ai/mcp/attic}/SSEClient.java | 2 +- pom.xml | 2 + samples/spring-ai-mcp-webmvc-sample/pom.xml | 7 + .../webmvc/client/ClientWebFluxSse.java | 32 ++ ...tion.java => McpMvcServerApplication.java} | 4 +- .../sample/webmvc/server/McpServerConfig.java | 10 - .../src/main/resources/application.properties | 4 +- 13 files changed, 726 insertions(+), 104 deletions(-) rename mcp-transport/mcp-webflux-sse-transport/src/test/java/org/springframework/ai/mcp/{server/WebFluxSseAsyncIntegrationTests.java => WebFluxSseIntegrationTests.java} (84%) create mode 100644 mcp-transport/mcp-webmvc-sse-transport/src/test/java/org/springframework/ai/mcp/server/WebMvcSseIntegrationTests.java rename mcp/src/{main/java/org/springframework/ai/mcp/client/transport => test/java/org/springframework/ai/mcp/attic}/SSEClient.java (98%) create mode 100644 samples/spring-ai-mcp-webmvc-sample/src/main/java/org/springframework/ai/mcp/sample/webmvc/client/ClientWebFluxSse.java rename samples/spring-ai-mcp-webmvc-sample/src/main/java/org/springframework/ai/mcp/sample/webmvc/server/{McpServerApplication.java => McpMvcServerApplication.java} (71%) diff --git a/mcp-transport/mcp-webflux-sse-transport/pom.xml b/mcp-transport/mcp-webflux-sse-transport/pom.xml index a4150e2..1dae71c 100644 --- a/mcp-transport/mcp-webflux-sse-transport/pom.xml +++ b/mcp-transport/mcp-webflux-sse-transport/pom.xml @@ -108,7 +108,14 @@ test + + org.junit.jupiter + junit-jupiter-params + ${junit-jupiter.version} + test + + - \ No newline at end of file + diff --git a/mcp-transport/mcp-webflux-sse-transport/src/test/java/org/springframework/ai/mcp/server/WebFluxSseAsyncIntegrationTests.java b/mcp-transport/mcp-webflux-sse-transport/src/test/java/org/springframework/ai/mcp/WebFluxSseIntegrationTests.java similarity index 84% rename from mcp-transport/mcp-webflux-sse-transport/src/test/java/org/springframework/ai/mcp/server/WebFluxSseAsyncIntegrationTests.java rename to mcp-transport/mcp-webflux-sse-transport/src/test/java/org/springframework/ai/mcp/WebFluxSseIntegrationTests.java index d9d601f..1e3c11a 100644 --- a/mcp-transport/mcp-webflux-sse-transport/src/test/java/org/springframework/ai/mcp/server/WebFluxSseAsyncIntegrationTests.java +++ b/mcp-transport/mcp-webflux-sse-transport/src/test/java/org/springframework/ai/mcp/WebFluxSseIntegrationTests.java @@ -13,11 +13,12 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.springframework.ai.mcp.server; +package org.springframework.ai.mcp; import java.time.Duration; import java.util.List; import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Function; @@ -25,12 +26,16 @@ import com.fasterxml.jackson.databind.ObjectMapper; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; import reactor.netty.DisposableServer; import reactor.netty.http.server.HttpServer; import reactor.test.StepVerifier; import org.springframework.ai.mcp.client.McpClient; +import org.springframework.ai.mcp.client.transport.HttpClientSseClientTransport; import org.springframework.ai.mcp.client.transport.WebFluxSseClientTransport; +import org.springframework.ai.mcp.server.McpServer; import org.springframework.ai.mcp.server.McpServer.ToolRegistration; import org.springframework.ai.mcp.server.transport.WebFluxSseServerTransport; import org.springframework.ai.mcp.spec.McpError; @@ -54,7 +59,7 @@ import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.awaitility.Awaitility.await; -public class WebFluxSseAsyncIntegrationTests { +public class WebFluxSseIntegrationTests { private static final int PORT = 8182; @@ -64,18 +69,21 @@ public class WebFluxSseAsyncIntegrationTests { private WebFluxSseServerTransport mcpServerTransport; - McpClient.Builder clientBuilder; + ConcurrentHashMap clientBulders = new ConcurrentHashMap<>(); @BeforeEach public void before() { + this.mcpServerTransport = new WebFluxSseServerTransport(new ObjectMapper(), MESSAGE_ENDPOINT); HttpHandler httpHandler = RouterFunctions.toHttpHandler(mcpServerTransport.getRouterFunction()); ReactorHttpHandlerAdapter adapter = new ReactorHttpHandlerAdapter(httpHandler); this.httpServer = HttpServer.create().port(PORT).handle(adapter).bindNow(); - this.clientBuilder = McpClient - .using(new WebFluxSseClientTransport(WebClient.builder().baseUrl("http://localhost:" + PORT))); + clientBulders.put("httpclient", McpClient.using(new HttpClientSseClientTransport("http://localhost:" + PORT))); + clientBulders.put("webflux", McpClient + .using(new WebFluxSseClientTransport(WebClient.builder().baseUrl("http://localhost:" + PORT)))); + } @AfterEach @@ -105,11 +113,14 @@ public class WebFluxSseAsyncIntegrationTests { }); } - @Test - void testCreateMessageWithoutSamplingCapabilities() { + @ParameterizedTest(name = "{0} : {displayName} ") + @ValueSource(strings = { "httpclient", "webflux" }) + void testCreateMessageWithoutSamplingCapabilities(String clientType) { var mcpAsyncServer = McpServer.using(mcpServerTransport).serverInfo("test-server", "1.0.0").async(); + var clientBuilder = clientBulders.get(clientType); + var client = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")).sync(); InitializeResult initResult = client.initialize(); @@ -128,8 +139,11 @@ public class WebFluxSseAsyncIntegrationTests { }); } - @Test - void testCreateMessageSuccess() throws InterruptedException { + @ParameterizedTest(name = "{0} : {displayName} ") + @ValueSource(strings = { "httpclient", "webflux" }) + void testCreateMessageSuccess(String clientType) throws InterruptedException { + + var clientBuilder = clientBulders.get(clientType); var mcpAsyncServer = McpServer.using(mcpServerTransport).serverInfo("test-server", "1.0.0").async(); @@ -169,8 +183,11 @@ public class WebFluxSseAsyncIntegrationTests { // --------------------------------------- // Roots Tests // --------------------------------------- - @Test - void testRootsSuccess() { + @ParameterizedTest(name = "{0} : {displayName} ") + @ValueSource(strings = { "httpclient", "webflux" }) + void testRootsSuccess(String clientType) { + var clientBuilder = clientBulders.get(clientType); + List roots = List.of(new Root("uri1://", "root1"), new Root("uri2://", "root2")); AtomicReference> rootsRef = new AtomicReference<>(); @@ -214,8 +231,11 @@ public class WebFluxSseAsyncIntegrationTests { mcpServer.close(); } - @Test - void testRootsWithoutCapability() { + @ParameterizedTest(name = "{0} : {displayName} ") + @ValueSource(strings = { "httpclient", "webflux" }) + void testRootsWithoutCapability(String clientType) { + var clientBuilder = clientBulders.get(clientType); + var mcpServer = McpServer.using(mcpServerTransport).rootsChangeConsumer(rootsUpdate -> { }).sync(); @@ -236,8 +256,11 @@ public class WebFluxSseAsyncIntegrationTests { mcpServer.close(); } - @Test - void testRootsWithEmptyRootsList() { + @ParameterizedTest(name = "{0} : {displayName} ") + @ValueSource(strings = { "httpclient", "webflux" }) + void testRootsWithEmptyRootsList(String clientType) { + var clientBuilder = clientBulders.get(clientType); + AtomicReference> rootsRef = new AtomicReference<>(); var mcpServer = McpServer.using(mcpServerTransport) .rootsChangeConsumer(rootsUpdate -> rootsRef.set(rootsUpdate)) @@ -260,8 +283,11 @@ public class WebFluxSseAsyncIntegrationTests { mcpServer.close(); } - @Test - void testRootsWithMultipleConsumers() { + @ParameterizedTest(name = "{0} : {displayName} ") + @ValueSource(strings = { "httpclient", "webflux" }) + void testRootsWithMultipleConsumers(String clientType) { + var clientBuilder = clientBulders.get(clientType); + List roots = List.of(new Root("uri1://", "root1")); AtomicReference> rootsRef1 = new AtomicReference<>(); @@ -290,8 +316,12 @@ public class WebFluxSseAsyncIntegrationTests { mcpServer.close(); } - @Test - void testRootsServerCloseWithActiveSubscription() { + @ParameterizedTest(name = "{0} : {displayName} ") + @ValueSource(strings = { "httpclient", "webflux" }) + void testRootsServerCloseWithActiveSubscription(String clientType) { + + var clientBuilder = clientBulders.get(clientType); + List roots = List.of(new Root("uri1://", "root1")); AtomicReference> rootsRef = new AtomicReference<>(); @@ -331,9 +361,11 @@ public class WebFluxSseAsyncIntegrationTests { } """; - @Test + @ParameterizedTest(name = "{0} : {displayName} ") + @ValueSource(strings = { "httpclient", "webflux" }) + void testToolCallSuccess(String clientType) { - void testToolCallSuccess() { + var clientBuilder = clientBulders.get(clientType); var callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), null); ToolRegistration tool1 = new ToolRegistration(new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema), @@ -369,8 +401,11 @@ public class WebFluxSseAsyncIntegrationTests { mcpServer.close(); } - @Test - void testToolListChangeHandlingSuccess() { + @ParameterizedTest(name = "{0} : {displayName} ") + @ValueSource(strings = { "httpclient", "webflux" }) + void testToolListChangeHandlingSuccess(String clientType) { + + var clientBuilder = clientBulders.get(clientType); var callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), null); ToolRegistration tool1 = new ToolRegistration(new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema), diff --git a/mcp-transport/mcp-webmvc-sse-transport/src/main/java/org/springframework/ai/mcp/server/transport/WebMvcSseServerTransport.java b/mcp-transport/mcp-webmvc-sse-transport/src/main/java/org/springframework/ai/mcp/server/transport/WebMvcSseServerTransport.java index 3305da5..7e675a7 100644 --- a/mcp-transport/mcp-webmvc-sse-transport/src/main/java/org/springframework/ai/mcp/server/transport/WebMvcSseServerTransport.java +++ b/mcp-transport/mcp-webmvc-sse-transport/src/main/java/org/springframework/ai/mcp/server/transport/WebMvcSseServerTransport.java @@ -1,5 +1,29 @@ +/* + * Copyright 2024-2024 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + package org.springframework.ai.mcp.server.transport; +import java.io.IOException; +import java.util.UUID; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.TimeUnit; +import java.util.function.Function; + import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; import org.slf4j.Logger; @@ -11,18 +35,11 @@ import org.springframework.ai.mcp.spec.McpSchema; import org.springframework.ai.mcp.spec.ServerMcpTransport; import org.springframework.ai.mcp.util.Assert; import org.springframework.http.HttpStatus; -import org.springframework.http.MediaType; -import org.springframework.web.servlet.mvc.method.annotation.SseEmitter; import org.springframework.web.servlet.function.RouterFunction; import org.springframework.web.servlet.function.RouterFunctions; import org.springframework.web.servlet.function.ServerRequest; import org.springframework.web.servlet.function.ServerResponse; -import java.io.IOException; -import java.util.UUID; -import java.util.concurrent.ConcurrentHashMap; -import java.util.function.Function; - /** * Server-side implementation of the MCP HTTP with SSE transport specification using * Spring WebMVC. This implementation wraps synchronous WebMVC operations in reactive @@ -105,9 +122,10 @@ public class WebMvcSseServerTransport implements ServerMcpTransport { sessions.values().forEach(session -> { try { - session.messageSink.send(SseEmitter.event().name(MESSAGE_EVENT_TYPE).data(jsonText)); + var event = new SSEEvent(session.id, MESSAGE_EVENT_TYPE, jsonText); + session.queue.put(event); } - catch (IOException e) { + catch (Exception e) { logger.error("Failed to send message to session {}: {}", session.id, e.getMessage()); } }); @@ -118,39 +136,11 @@ public class WebMvcSseServerTransport implements ServerMcpTransport { }); } - @Override - public T unmarshalFrom(Object data, TypeReference typeRef) { - return this.objectMapper.convertValue(data, typeRef); - } - - @Override - public Mono closeGracefully() { - return Mono.fromRunnable(() -> { - isClosing = true; - logger.debug("Initiating graceful shutdown with {} active sessions", sessions.size()); - - sessions.values().forEach(session -> { - String sessionId = session.id; - session.close(); - sessions.remove(sessionId); - }); - - logger.info("Graceful shutdown completed"); - }); - } - - /** - * Get the router function for configuring the web server. - */ - public RouterFunction getRouterFunction() { - return this.routerFunction; - } - /** * Handles new SSE connection requests from clients. */ private ServerResponse handleSseConnection(ServerRequest request) { - if (isClosing) { + if (this.isClosing) { return ServerResponse.status(HttpStatus.SERVICE_UNAVAILABLE).body("Server is shutting down"); } @@ -162,37 +152,36 @@ public class WebMvcSseServerTransport implements ServerMcpTransport { // Send initial endpoint event try { - session.messageSink.send(SseEmitter.event().name(ENDPOINT_EVENT_TYPE).data(messageEndpoint)); + session.queue.put(new SSEEvent(session.id, ENDPOINT_EVENT_TYPE, messageEndpoint)); + return ServerResponse.sse(sseBuilder -> { + // new Thread(() -> { + while (!this.isClosing) { + try { + SSEEvent sseEvent = session.queue.poll(100, TimeUnit.MILLISECONDS); + if (sseEvent != null) { + sseBuilder.id(sseEvent.id).event(sseEvent.type()).data(sseEvent.data()); + } + } + catch (Exception e) { + logger.error("Failed to poll event from session queue: {}", e.getMessage()); + sseBuilder.error(e); + } + } + // }).start(); + }); } - catch (IOException e) { + catch (Exception e) { logger.error("Failed to send initial endpoint event to session {}: {}", sessionId, e.getMessage()); sessions.remove(sessionId); return ServerResponse.status(HttpStatus.INTERNAL_SERVER_ERROR).build(); } - - session.messageSink.onCompletion(() -> { - logger.debug("Session {} completed", sessionId); - sessions.remove(sessionId); - }); - - session.messageSink.onError(error -> { - logger.error("Error in session {}: {}", sessionId, error.getMessage()); - sessions.remove(sessionId); - }); - - session.messageSink.onTimeout(() -> { - logger.debug("Session {} timed out", sessionId); - sessions.remove(sessionId); - }); - - return ServerResponse.ok().contentType(MediaType.TEXT_EVENT_STREAM).body(session.messageSink); } /** * Handles incoming messages from clients. */ private ServerResponse handleMessage(ServerRequest request) { - if (isClosing) { + if (this.isClosing) { return ServerResponse.status(HttpStatus.SERVICE_UNAVAILABLE).body("Server is shutting down"); } @@ -216,6 +205,9 @@ public class WebMvcSseServerTransport implements ServerMcpTransport { } } + record SSEEvent(String id, String type, String data) { + } + /** * Represents an active client session. */ @@ -223,18 +215,18 @@ public class WebMvcSseServerTransport implements ServerMcpTransport { private final String id; - private final SseEmitter messageSink; + private final BlockingQueue queue; ClientSession(String id) { this.id = id; - this.messageSink = new SseEmitter(0L); // No timeout + this.queue = new LinkedBlockingQueue<>(); logger.debug("Session {} initialized with SSE emitter", id); } void close() { logger.debug("Closing session: {}", id); try { - messageSink.complete(); + queue.remove(); logger.debug("Successfully completed SSE emitter for session {}", id); } catch (Exception e) { @@ -244,4 +236,32 @@ public class WebMvcSseServerTransport implements ServerMcpTransport { } + @Override + public T unmarshalFrom(Object data, TypeReference typeRef) { + return this.objectMapper.convertValue(data, typeRef); + } + + @Override + public Mono closeGracefully() { + return Mono.fromRunnable(() -> { + this.isClosing = true; + logger.debug("Initiating graceful shutdown with {} active sessions", sessions.size()); + + sessions.values().forEach(session -> { + String sessionId = session.id; + session.close(); + sessions.remove(sessionId); + }); + + logger.info("Graceful shutdown completed"); + }); + } + + /** + * Get the router function for configuring the web server. + */ + public RouterFunction getRouterFunction() { + return this.routerFunction; + } + } diff --git a/mcp-transport/mcp-webmvc-sse-transport/src/test/java/org/springframework/ai/mcp/server/WebMvcSseIntegrationTests.java b/mcp-transport/mcp-webmvc-sse-transport/src/test/java/org/springframework/ai/mcp/server/WebMvcSseIntegrationTests.java new file mode 100644 index 0000000..6658386 --- /dev/null +++ b/mcp-transport/mcp-webmvc-sse-transport/src/test/java/org/springframework/ai/mcp/server/WebMvcSseIntegrationTests.java @@ -0,0 +1,524 @@ +/* + * Copyright 2024 - 2024 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.springframework.ai.mcp.server; + +import java.time.Duration; +import java.util.List; +import java.util.Map; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Function; + +import com.fasterxml.jackson.databind.ObjectMapper; +import org.apache.catalina.Context; +import org.apache.catalina.LifecycleException; +import org.apache.catalina.LifecycleState; +import org.apache.catalina.startup.Tomcat; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import reactor.test.StepVerifier; + +import org.springframework.ai.mcp.client.McpClient; +import org.springframework.ai.mcp.client.transport.HttpClientSseClientTransport; +import org.springframework.ai.mcp.server.McpServer.ToolRegistration; +import org.springframework.ai.mcp.server.transport.WebMvcSseServerTransport; +import org.springframework.ai.mcp.spec.McpError; +import org.springframework.ai.mcp.spec.McpSchema; +import org.springframework.ai.mcp.spec.McpSchema.CallToolResult; +import org.springframework.ai.mcp.spec.McpSchema.ClientCapabilities; +import org.springframework.ai.mcp.spec.McpSchema.CreateMessageRequest; +import org.springframework.ai.mcp.spec.McpSchema.CreateMessageResult; +import org.springframework.ai.mcp.spec.McpSchema.InitializeResult; +import org.springframework.ai.mcp.spec.McpSchema.Role; +import org.springframework.ai.mcp.spec.McpSchema.Root; +import org.springframework.ai.mcp.spec.McpSchema.ServerCapabilities; +import org.springframework.ai.mcp.spec.McpSchema.Tool; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Configuration; +import org.springframework.web.client.RestClient; +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 static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.awaitility.Awaitility.await; + +public class WebMvcSseIntegrationTests { + + private static final Logger logger = LoggerFactory.getLogger(WebMvcSseIntegrationTests.class); + + private static final int PORT = 8183; + + private static final String MESSAGE_ENDPOINT = "/mcp/message"; + + private WebMvcSseServerTransport mcpServerTransport; + + McpClient.Builder clientBuilder; + + @Configuration + @EnableWebMvc + static class TestConfig { + + @Bean + public WebMvcSseServerTransport webMvcSseServerTransport() { + return new WebMvcSseServerTransport(new ObjectMapper(), MESSAGE_ENDPOINT); + } + + @Bean + public RouterFunction routerFunction(WebMvcSseServerTransport transport) { + return transport.getRouterFunction(); + } + + } + + private Tomcat tomcat; + + private AnnotationConfigWebApplicationContext appContext; + + @BeforeEach + public void before() { + + // 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 + mcpServerTransport = appContext.getBean(WebMvcSseServerTransport.class); + + // Create DispatcherServlet with our Spring context + DispatcherServlet dispatcherServlet = new DispatcherServlet(appContext); + // dispatcherServlet.setThrowExceptionIfNoHandlerFound(true); + + // Add servlet to Tomcat and get the wrapper + var wrapper = Tomcat.addServlet(context, "dispatcherServlet", dispatcherServlet); + wrapper.setLoadOnStartup(1); + wrapper.setAsyncSupported(true); + context.addServletMappingDecoded("/*", "dispatcherServlet"); + + try { + // Configure and start the connector with async support + var connector = tomcat.getConnector(); + connector.setAsyncTimeout(3000); // 3 seconds timeout for async requests + tomcat.start(); + assertThat(tomcat.getServer().getState() == LifecycleState.STARTED); + } + catch (Exception e) { + throw new RuntimeException("Failed to start Tomcat", e); + } + + this.clientBuilder = McpClient.using(new HttpClientSseClientTransport("http://localhost:" + PORT)); + } + + @AfterEach + public void after() { + if (mcpServerTransport != null) { + mcpServerTransport.closeGracefully().block(); + } + if (appContext != null) { + appContext.close(); + } + if (tomcat != null) { + try { + tomcat.stop(); + tomcat.destroy(); + } + catch (LifecycleException e) { + throw new RuntimeException("Failed to stop Tomcat", e); + } + } + } + + // --------------------------------------- + // Sampling Tests + // --------------------------------------- + @Test + void testCreateMessageWithoutInitialization() { + var mcpAsyncServer = McpServer.using(mcpServerTransport).serverInfo("test-server", "1.0.0").async(); + + var messages = List + .of(new McpSchema.SamplingMessage(McpSchema.Role.USER, new McpSchema.TextContent("Test message"))); + var modelPrefs = new McpSchema.ModelPreferences(List.of(), 1.0, 1.0, 1.0); + + var request = new McpSchema.CreateMessageRequest(messages, modelPrefs, null, + McpSchema.CreateMessageRequest.ContextInclusionStrategy.NONE, null, 100, List.of(), Map.of()); + + StepVerifier.create(mcpAsyncServer.createMessage(request)).verifyErrorSatisfies(error -> { + assertThat(error).isInstanceOf(McpError.class) + .hasMessage("Client must be initialized. Call the initialize method first!"); + }); + } + + @Test + void testCreateMessageWithoutSamplingCapabilities() { + + var mcpAsyncServer = McpServer.using(mcpServerTransport).serverInfo("test-server", "1.0.0").async(); + + var client = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")).sync(); + + InitializeResult initResult = client.initialize(); + assertThat(initResult).isNotNull(); + + var messages = List + .of(new McpSchema.SamplingMessage(McpSchema.Role.USER, new McpSchema.TextContent("Test message"))); + var modelPrefs = new McpSchema.ModelPreferences(List.of(), 1.0, 1.0, 1.0); + + var request = new McpSchema.CreateMessageRequest(messages, modelPrefs, null, + McpSchema.CreateMessageRequest.ContextInclusionStrategy.NONE, null, 100, List.of(), Map.of()); + + StepVerifier.create(mcpAsyncServer.createMessage(request)).verifyErrorSatisfies(error -> { + assertThat(error).isInstanceOf(McpError.class) + .hasMessage("Client must be configured with sampling capabilities"); + }); + } + + @Test + void testCreateMessageSuccess() throws InterruptedException { + + var mcpAsyncServer = McpServer.using(mcpServerTransport).serverInfo("test-server", "1.0.0").async(); + + Function samplingHandler = request -> { + assertThat(request.messages()).hasSize(1); + assertThat(request.messages().get(0).content()).isInstanceOf(McpSchema.TextContent.class); + + return new CreateMessageResult(Role.USER, new McpSchema.TextContent("Test message"), "MockModelName", + CreateMessageResult.StopReason.STOP_SEQUENCE); + }; + + var client = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")) + .capabilities(ClientCapabilities.builder().sampling().build()) + .sampling(samplingHandler) + .sync(); + + InitializeResult initResult = client.initialize(); + assertThat(initResult).isNotNull(); + + var messages = List + .of(new McpSchema.SamplingMessage(McpSchema.Role.USER, new McpSchema.TextContent("Test message"))); + var modelPrefs = new McpSchema.ModelPreferences(List.of(), 1.0, 1.0, 1.0); + + var request = new McpSchema.CreateMessageRequest(messages, modelPrefs, null, + McpSchema.CreateMessageRequest.ContextInclusionStrategy.NONE, null, 100, List.of(), Map.of()); + + StepVerifier.create(mcpAsyncServer.createMessage(request)).consumeNextWith(result -> { + assertThat(result).isNotNull(); + assertThat(result.role()).isEqualTo(Role.USER); + assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class); + assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message"); + assertThat(result.model()).isEqualTo("MockModelName"); + assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE); + }).verifyComplete(); + } + + // --------------------------------------- + // Roots Tests + // --------------------------------------- + @Test + void testRootsSuccess() { + List roots = List.of(new Root("uri1://", "root1"), new Root("uri2://", "root2")); + + AtomicReference> rootsRef = new AtomicReference<>(); + var mcpServer = McpServer.using(mcpServerTransport) + .rootsChangeConsumer(rootsUpdate -> rootsRef.set(rootsUpdate)) + .sync(); + + var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().roots(true).build()) + .roots(roots) + .sync(); + + InitializeResult initResult = mcpClient.initialize(); + assertThat(initResult).isNotNull(); + + assertThat(rootsRef.get()).isNull(); + + assertThat(mcpServer.listRoots().roots()).containsAll(roots); + + mcpClient.rootsListChangedNotification(); + + await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { + assertThat(rootsRef.get()).containsAll(roots); + }); + + // Remove a root + mcpClient.removeRoot(roots.get(0).uri()); + + await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { + assertThat(rootsRef.get()).containsAll(List.of(roots.get(1))); + }); + + // Add a new root + var root3 = new Root("uri3://", "root3"); + mcpClient.addRoot(root3); + + await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { + assertThat(rootsRef.get()).containsAll(List.of(roots.get(1), root3)); + }); + + mcpClient.close(); + mcpServer.close(); + } + + @Test + void testRootsWithoutCapability() { + var mcpServer = McpServer.using(mcpServerTransport).rootsChangeConsumer(rootsUpdate -> { + }).sync(); + + // Create client without roots capability + // No roots capability + var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().build()).sync(); + + InitializeResult initResult = mcpClient.initialize(); + assertThat(initResult).isNotNull(); + + // Attempt to list roots should fail + assertThatThrownBy(() -> mcpServer.listRoots().roots()).isInstanceOf(McpError.class) + .hasMessage("Roots not supported"); + + mcpClient.close(); + mcpServer.close(); + } + + @Test + void testRootsWithEmptyRootsList() { + AtomicReference> rootsRef = new AtomicReference<>(); + var mcpServer = McpServer.using(mcpServerTransport) + .rootsChangeConsumer(rootsUpdate -> rootsRef.set(rootsUpdate)) + .sync(); + + var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().roots(true).build()) + .roots(List.of()) // Empty roots list + .sync(); + + InitializeResult initResult = mcpClient.initialize(); + assertThat(initResult).isNotNull(); + + mcpClient.rootsListChangedNotification(); + + await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { + assertThat(rootsRef.get()).isEmpty(); + }); + + mcpClient.close(); + mcpServer.close(); + } + + @Test + void testRootsWithMultipleConsumers() { + List roots = List.of(new Root("uri1://", "root1")); + + AtomicReference> rootsRef1 = new AtomicReference<>(); + AtomicReference> rootsRef2 = new AtomicReference<>(); + + var mcpServer = McpServer.using(mcpServerTransport) + .rootsChangeConsumer(rootsUpdate -> rootsRef1.set(rootsUpdate)) + .rootsChangeConsumer(rootsUpdate -> rootsRef2.set(rootsUpdate)) + .sync(); + + var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().roots(true).build()) + .roots(roots) + .sync(); + + InitializeResult initResult = mcpClient.initialize(); + assertThat(initResult).isNotNull(); + + mcpClient.rootsListChangedNotification(); + + await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { + assertThat(rootsRef1.get()).containsAll(roots); + assertThat(rootsRef2.get()).containsAll(roots); + }); + + mcpClient.close(); + mcpServer.close(); + } + + @Test + void testRootsServerCloseWithActiveSubscription() { + List roots = List.of(new Root("uri1://", "root1")); + + AtomicReference> rootsRef = new AtomicReference<>(); + var mcpServer = McpServer.using(mcpServerTransport) + .rootsChangeConsumer(rootsUpdate -> rootsRef.set(rootsUpdate)) + .sync(); + + var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().roots(true).build()) + .roots(roots) + .sync(); + + InitializeResult initResult = mcpClient.initialize(); + assertThat(initResult).isNotNull(); + + mcpClient.rootsListChangedNotification(); + + await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { + assertThat(rootsRef.get()).containsAll(roots); + }); + + // Close server while subscription is active + mcpServer.close(); + + // Verify client can handle server closure gracefully + mcpClient.close(); + } + + // --------------------------------------- + // Tools Tests + // --------------------------------------- + + String emptyJsonSchema = """ + { + "$schema": "http://json-schema.org/draft-07/schema#", + "type": "object", + "properties": {} + } + """; + + @Test + void testToolCallSuccess() { + + var callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), null); + ToolRegistration tool1 = new ToolRegistration(new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema), + request -> { + // perform a blocking call to a remote service + String response = RestClient.create() + .get() + .uri("https://github.com/spring-projects-experimental/spring-ai-mcp/blob/main/README.md") + .retrieve() + .body(String.class); + assertThat(response).isNotBlank(); + return callResponse; + }); + + var mcpServer = McpServer.using(mcpServerTransport) + .capabilities(ServerCapabilities.builder().tools(true).build()) + .tools(tool1) + .sync(); + + var mcpClient = clientBuilder.sync(); + + InitializeResult initResult = mcpClient.initialize(); + assertThat(initResult).isNotNull(); + + assertThat(mcpClient.listTools().tools()).contains(tool1.tool()); + + CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); + + assertThat(response).isNotNull(); + assertThat(response).isEqualTo(callResponse); + + mcpClient.close(); + mcpServer.close(); + } + + @Test + void testToolListChangeHandlingSuccess() { + + var callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), null); + ToolRegistration tool1 = new ToolRegistration(new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema), + request -> { + // perform a blocking call to a remote service + String response = RestClient.create() + .get() + .uri("https://github.com/spring-projects-experimental/spring-ai-mcp/blob/main/README.md") + .retrieve() + .body(String.class); + assertThat(response).isNotBlank(); + return callResponse; + }); + + var mcpServer = McpServer.using(mcpServerTransport) + .capabilities(ServerCapabilities.builder().tools(true).build()) + .tools(tool1) + .sync(); + + AtomicReference> rootsRef = new AtomicReference<>(); + var mcpClient = clientBuilder.toolsChangeConsumer(toolsUpdate -> { + // perform a blocking call to a remote service + String response = RestClient.create() + .get() + .uri("https://github.com/spring-projects-experimental/spring-ai-mcp/blob/main/README.md") + .retrieve() + .body(String.class); + assertThat(response).isNotBlank(); + rootsRef.set(toolsUpdate); + }).sync(); + + InitializeResult initResult = mcpClient.initialize(); + assertThat(initResult).isNotNull(); + + assertThat(rootsRef.get()).isNull(); + + assertThat(mcpClient.listTools().tools()).contains(tool1.tool()); + + mcpServer.notifyToolsListChanged(); + + await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { + assertThat(rootsRef.get()).containsAll(List.of(tool1.tool())); + }); + + // Remove a tool + mcpServer.removeTool("tool1"); + + await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { + assertThat(rootsRef.get()).isEmpty(); + }); + + // Add a new tool + ToolRegistration tool2 = new ToolRegistration(new McpSchema.Tool("tool2", "tool2 description", emptyJsonSchema), + request -> callResponse); + + mcpServer.addTool(tool2); + + await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { + assertThat(rootsRef.get()).containsAll(List.of(tool2.tool())); + }); + + mcpClient.close(); + mcpServer.close(); + } + + @Test + void testInitialize() { + + var mcpServer = McpServer.using(mcpServerTransport).sync(); + + var mcpClient = clientBuilder.sync(); + + InitializeResult initResult = mcpClient.initialize(); + assertThat(initResult).isNotNull(); + + mcpClient.close(); + mcpServer.close(); + } + +} diff --git a/mcp-transport/mcp-webmvc-sse-transport/src/test/resources/logback.xml b/mcp-transport/mcp-webmvc-sse-transport/src/test/resources/logback.xml index 0246d6c..517af52 100644 --- a/mcp-transport/mcp-webmvc-sse-transport/src/test/resources/logback.xml +++ b/mcp-transport/mcp-webmvc-sse-transport/src/test/resources/logback.xml @@ -9,13 +9,16 @@ - + - + + + + - + diff --git a/mcp/src/main/java/org/springframework/ai/mcp/client/transport/HttpClientSseClientTransport.java b/mcp/src/main/java/org/springframework/ai/mcp/client/transport/HttpClientSseClientTransport.java index 6423682..3f457a4 100644 --- a/mcp/src/main/java/org/springframework/ai/mcp/client/transport/HttpClientSseClientTransport.java +++ b/mcp/src/main/java/org/springframework/ai/mcp/client/transport/HttpClientSseClientTransport.java @@ -139,7 +139,9 @@ public class HttpClientSseClientTransport implements ClientMcpTransport { } try { - closeLatch.await(10, TimeUnit.SECONDS); + if (!closeLatch.await(10, TimeUnit.SECONDS)) { + return Mono.error(new McpError("Failed to wait for the message endpoint")); + } } catch (InterruptedException e) { return Mono.error(new McpError("Failed to wait for the message endpoint")); diff --git a/mcp/src/main/java/org/springframework/ai/mcp/client/transport/SSEClient.java b/mcp/src/test/java/org/springframework/ai/mcp/attic/SSEClient.java similarity index 98% rename from mcp/src/main/java/org/springframework/ai/mcp/client/transport/SSEClient.java rename to mcp/src/test/java/org/springframework/ai/mcp/attic/SSEClient.java index ec8c8eb..3b75999 100644 --- a/mcp/src/main/java/org/springframework/ai/mcp/client/transport/SSEClient.java +++ b/mcp/src/test/java/org/springframework/ai/mcp/attic/SSEClient.java @@ -13,7 +13,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.springframework.ai.mcp.client.transport; +package org.springframework.ai.mcp.attic; /** * @author Christian Tzolov diff --git a/pom.xml b/pom.xml index 1b63284..c906585 100644 --- a/pom.xml +++ b/pom.xml @@ -96,6 +96,7 @@ 1.0.0-alpha.4 0.0.4 1.6.2 + 5.10.5 @@ -105,6 +106,7 @@ spring-ai-mcp mcp-docs mcp-test + diff --git a/samples/spring-ai-mcp-webmvc-sample/pom.xml b/samples/spring-ai-mcp-webmvc-sample/pom.xml index f99d2b3..ba29714 100644 --- a/samples/spring-ai-mcp-webmvc-sample/pom.xml +++ b/samples/spring-ai-mcp-webmvc-sample/pom.xml @@ -27,6 +27,13 @@ mcp-webmvc-sse-transport ${spring-ai-mcp.version} + + + + org.springframework.experimental + mcp-webflux-sse-transport + ${spring-ai-mcp.version} + org.springframework.experimental spring-ai-mcp diff --git a/samples/spring-ai-mcp-webmvc-sample/src/main/java/org/springframework/ai/mcp/sample/webmvc/client/ClientWebFluxSse.java b/samples/spring-ai-mcp-webmvc-sample/src/main/java/org/springframework/ai/mcp/sample/webmvc/client/ClientWebFluxSse.java new file mode 100644 index 0000000..60f5609 --- /dev/null +++ b/samples/spring-ai-mcp-webmvc-sample/src/main/java/org/springframework/ai/mcp/sample/webmvc/client/ClientWebFluxSse.java @@ -0,0 +1,32 @@ +/* +* Copyright 2024 - 2024 the original author or authors. +* +* Licensed under the Apache License, Version 2.0 (the "License"); +* you may not use this file except in compliance with the License. +* You may obtain a copy of the License at +* +* https://www.apache.org/licenses/LICENSE-2.0 +* +* Unless required by applicable law or agreed to in writing, software +* distributed under the License is distributed on an "AS IS" BASIS, +* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +* See the License for the specific language governing permissions and +* limitations under the License. +*/ +package org.springframework.ai.mcp.sample.webmvc.client; + +import org.springframework.ai.mcp.client.transport.WebFluxSseClientTransport; +import org.springframework.web.reactive.function.client.WebClient; + +/** + * @author Christian Tzolov + */ +public class ClientWebFluxSse { + + public static void main(String[] args) { + var transport = new WebFluxSseClientTransport(WebClient.builder().baseUrl("http://localhost:8080")); + + new SampleClient(transport).run(); + } + +} diff --git a/samples/spring-ai-mcp-webmvc-sample/src/main/java/org/springframework/ai/mcp/sample/webmvc/server/McpServerApplication.java b/samples/spring-ai-mcp-webmvc-sample/src/main/java/org/springframework/ai/mcp/sample/webmvc/server/McpMvcServerApplication.java similarity index 71% rename from samples/spring-ai-mcp-webmvc-sample/src/main/java/org/springframework/ai/mcp/sample/webmvc/server/McpServerApplication.java rename to samples/spring-ai-mcp-webmvc-sample/src/main/java/org/springframework/ai/mcp/sample/webmvc/server/McpMvcServerApplication.java index 82f80de..cfff5e0 100644 --- a/samples/spring-ai-mcp-webmvc-sample/src/main/java/org/springframework/ai/mcp/sample/webmvc/server/McpServerApplication.java +++ b/samples/spring-ai-mcp-webmvc-sample/src/main/java/org/springframework/ai/mcp/sample/webmvc/server/McpMvcServerApplication.java @@ -4,10 +4,10 @@ import org.springframework.boot.SpringApplication; import org.springframework.boot.autoconfigure.SpringBootApplication; @SpringBootApplication -public class McpServerApplication { +public class McpMvcServerApplication { public static void main(String[] args) { - SpringApplication.run(McpServerApplication.class, args); + SpringApplication.run(McpMvcServerApplication.class, args); } } diff --git a/samples/spring-ai-mcp-webmvc-sample/src/main/java/org/springframework/ai/mcp/sample/webmvc/server/McpServerConfig.java b/samples/spring-ai-mcp-webmvc-sample/src/main/java/org/springframework/ai/mcp/sample/webmvc/server/McpServerConfig.java index be2ea57..573dc31 100644 --- a/samples/spring-ai-mcp-webmvc-sample/src/main/java/org/springframework/ai/mcp/sample/webmvc/server/McpServerConfig.java +++ b/samples/spring-ai-mcp-webmvc-sample/src/main/java/org/springframework/ai/mcp/sample/webmvc/server/McpServerConfig.java @@ -31,8 +31,6 @@ import org.springframework.context.annotation.Configuration; import org.springframework.web.client.RestClient; import org.springframework.web.servlet.config.annotation.EnableWebMvc; import org.springframework.web.servlet.config.annotation.WebMvcConfigurer; -import org.springframework.http.converter.HttpMessageConverter; -import org.springframework.http.converter.json.MappingJackson2HttpMessageConverter; import org.springframework.web.servlet.function.RouterFunction; import org.springframework.web.servlet.function.ServerResponse; @@ -40,14 +38,6 @@ import org.springframework.web.servlet.function.ServerResponse; @EnableWebMvc public class McpServerConfig implements WebMvcConfigurer { - @Override - public void configureMessageConverters(List> converters) { - MappingJackson2HttpMessageConverter converter = new MappingJackson2HttpMessageConverter(); - converter.setSupportedMediaTypes(List.of(org.springframework.http.MediaType.APPLICATION_JSON, - org.springframework.http.MediaType.TEXT_EVENT_STREAM)); - converters.add(converter); - } - private static final Logger logger = LoggerFactory.getLogger(McpServerConfig.class); @Bean diff --git a/samples/spring-ai-mcp-webmvc-sample/src/main/resources/application.properties b/samples/spring-ai-mcp-webmvc-sample/src/main/resources/application.properties index 2530c58..4bc4fbf 100644 --- a/samples/spring-ai-mcp-webmvc-sample/src/main/resources/application.properties +++ b/samples/spring-ai-mcp-webmvc-sample/src/main/resources/application.properties @@ -1,8 +1,8 @@ -# spring.main.web-application-type=none +spring.main.web-application-type=SERVLET # NOTE: You must disable the banner and the console logging # to allow the STDIO transport to work !!! spring.main.banner-mode=off -# logging.pattern.console= +logging.pattern.console= transport.mode=sse \ No newline at end of file