refactor: Split transport implementations into modules
Split the transport implementations into dedicated modules to improve modularity and reduce core dependencies. Added shared test module for common test utilities. Key changes: - Create mcp-transport/mcp-webflux-sse-transport module - Move SseClientTransport → WebFluxSseClientTransport - Move SseServerTransport → WebFluxSseServerTransport - Move related tests and utilities - Create mcp-transport/mcp-webmvc-sse-transport module - Add new WebMvcSseServerTransport implementation - Add corresponding tests and integration tests - Add framework-independent SSE client implementations to core module - Add FlowSseClient using java.net.http - Add HttpClientSseClientTransport using java.net.http - Keep StdioClientTransport/StdioServerTransport in core (no external deps) - Create mcp-test module for shared test infrastructure - Move abstract test base classes - Move shared test utilities and mocks - Update test dependencies The modular approach allows users to choose between: - Framework-specific implementations (WebFlux/WebMVC) - Pure Java implementations (java.net.http) - Basic implementations (STDIO) Add breaking changes readme: 0.5.0-BREAKING-CHANGES.md Resolves #54
This commit is contained in:
Christian Tzolov
parent
85c4fb8713
commit
611f3258d6
11 changed files with 5002 additions and 1588 deletions
|
|
@ -0,0 +1,61 @@
|
|||
/*
|
||||
* 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.attic;
|
||||
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import reactor.netty.http.server.HttpServer;
|
||||
|
||||
import org.springframework.ai.mcp.server.McpServer;
|
||||
import org.springframework.ai.mcp.server.transport.SseServerTransport;
|
||||
import org.springframework.ai.mcp.spec.McpSchema;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.CallToolResult;
|
||||
import org.springframework.http.server.reactive.HttpHandler;
|
||||
import org.springframework.http.server.reactive.ReactorHttpHandlerAdapter;
|
||||
import org.springframework.web.reactive.function.server.RouterFunctions;
|
||||
|
||||
/**
|
||||
*
|
||||
* https://docs.spring.io/spring-framework/reference/web/webflux-functional.html#webflux-fn-running
|
||||
* https://docs.spring.io/spring-framework/reference/web/webflux/reactive-spring.html#webflux-httphandler
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
*/
|
||||
|
||||
public class DemoServer {
|
||||
|
||||
public static void main(String[] args) {
|
||||
SseServerTransport transport = new SseServerTransport(new ObjectMapper(), "/mcp/message");
|
||||
|
||||
var mcpServer = McpServer.using(transport)
|
||||
.serverInfo("Weather Forecast", "1.0.0")
|
||||
.tool(new McpSchema.Tool("weather", "Weather forecast tool by location", "{ \"type\": \"object\" }"),
|
||||
(arguments) -> {
|
||||
String city = (String) arguments.get("city");
|
||||
return new CallToolResult(List.of(), false);
|
||||
})
|
||||
.async();
|
||||
|
||||
HttpHandler httpHandler = RouterFunctions.toHttpHandler(transport.getRouterFunction());
|
||||
ReactorHttpHandlerAdapter adapter = new ReactorHttpHandlerAdapter(httpHandler);
|
||||
HttpServer httpServer = HttpServer.create().port(8080).handle(adapter);
|
||||
httpServer.bindNow().onDispose().block();
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,297 @@
|
|||
/*
|
||||
* 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.BufferedReader;
|
||||
import java.io.IOException;
|
||||
import java.io.PrintWriter;
|
||||
import java.util.Map;
|
||||
import java.util.UUID;
|
||||
import java.util.concurrent.ConcurrentHashMap;
|
||||
import java.util.concurrent.atomic.AtomicBoolean;
|
||||
import java.util.function.Function;
|
||||
|
||||
import com.fasterxml.jackson.core.type.TypeReference;
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import jakarta.servlet.AsyncContext;
|
||||
import jakarta.servlet.ServletException;
|
||||
import jakarta.servlet.annotation.WebServlet;
|
||||
import jakarta.servlet.http.HttpServlet;
|
||||
import jakarta.servlet.http.HttpServletRequest;
|
||||
import jakarta.servlet.http.HttpServletResponse;
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
import reactor.core.publisher.Mono;
|
||||
|
||||
import org.springframework.ai.mcp.spec.McpError;
|
||||
import org.springframework.ai.mcp.spec.McpSchema;
|
||||
import org.springframework.ai.mcp.spec.ServerMcpTransport;
|
||||
|
||||
/**
|
||||
* A Servlet-based implementation of the MCP HTTP with SSE transport specification. This
|
||||
* implementation provides similar functionality to FunctionalSseServerTransport but uses
|
||||
* the traditional Servlet API instead of WebFlux.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
* @since 1.0.0
|
||||
*/
|
||||
|
||||
@WebServlet(asyncSupported = true)
|
||||
public class HttpServletSseServerTransport extends HttpServlet implements ServerMcpTransport {
|
||||
|
||||
private static final Logger logger = LoggerFactory.getLogger(HttpServletSseServerTransport.class);
|
||||
|
||||
public static final String SSE_ENDPOINT = "/sse";
|
||||
|
||||
public static final String MESSAGE_EVENT_TYPE = "message";
|
||||
|
||||
public static final String ENDPOINT_EVENT_TYPE = "endpoint";
|
||||
|
||||
private final ObjectMapper objectMapper;
|
||||
|
||||
private final String messageEndpoint;
|
||||
|
||||
private final Map<String, ClientSession> sessions = new ConcurrentHashMap<>();
|
||||
|
||||
private final AtomicBoolean isClosing = new AtomicBoolean(false);
|
||||
|
||||
private Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> connectHandler;
|
||||
|
||||
public HttpServletSseServerTransport(ObjectMapper objectMapper, String messageEndpoint) {
|
||||
this.objectMapper = objectMapper;
|
||||
this.messageEndpoint = messageEndpoint;
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void doGet(HttpServletRequest request, HttpServletResponse response)
|
||||
throws ServletException, IOException {
|
||||
|
||||
String pathInfo = request.getPathInfo();
|
||||
if (!SSE_ENDPOINT.equals(pathInfo)) {
|
||||
response.sendError(HttpServletResponse.SC_NOT_FOUND);
|
||||
return;
|
||||
}
|
||||
|
||||
if (isClosing.get()) {
|
||||
response.sendError(HttpServletResponse.SC_SERVICE_UNAVAILABLE, "Server is shutting down");
|
||||
return;
|
||||
}
|
||||
|
||||
response.setContentType("text/event-stream");
|
||||
response.setCharacterEncoding("UTF-8");
|
||||
response.setHeader("Cache-Control", "no-cache");
|
||||
response.setHeader("Connection", "keep-alive");
|
||||
response.setHeader("Access-Control-Allow-Origin", "*");
|
||||
|
||||
String sessionId = UUID.randomUUID().toString();
|
||||
AsyncContext asyncContext = request.startAsync();
|
||||
asyncContext.setTimeout(0);
|
||||
|
||||
PrintWriter writer = response.getWriter();
|
||||
ClientSession session = new ClientSession(sessionId, asyncContext, writer);
|
||||
sessions.put(sessionId, session);
|
||||
|
||||
// Send initial endpoint event
|
||||
sendEvent(writer, ENDPOINT_EVENT_TYPE, messageEndpoint);
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void doPost(HttpServletRequest request, HttpServletResponse response)
|
||||
throws ServletException, IOException {
|
||||
|
||||
if (isClosing.get()) {
|
||||
response.sendError(HttpServletResponse.SC_SERVICE_UNAVAILABLE, "Server is shutting down");
|
||||
return;
|
||||
}
|
||||
|
||||
String pathInfo = request.getPathInfo();
|
||||
if (!messageEndpoint.equals(pathInfo)) {
|
||||
response.sendError(HttpServletResponse.SC_NOT_FOUND);
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
BufferedReader reader = request.getReader();
|
||||
StringBuilder body = new StringBuilder();
|
||||
String line;
|
||||
while ((line = reader.readLine()) != null) {
|
||||
body.append(line);
|
||||
}
|
||||
|
||||
McpSchema.JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(objectMapper, body.toString());
|
||||
|
||||
if (connectHandler != null) {
|
||||
connectHandler.apply(Mono.just(message)).subscribe(responseMessage -> {
|
||||
try {
|
||||
response.setContentType("application/json");
|
||||
response.setCharacterEncoding("UTF-8");
|
||||
String jsonResponse = objectMapper.writeValueAsString(responseMessage);
|
||||
PrintWriter writer = response.getWriter();
|
||||
writer.write(jsonResponse);
|
||||
writer.flush();
|
||||
}
|
||||
catch (Exception e) {
|
||||
logger.error("Error sending response: {}", e.getMessage());
|
||||
try {
|
||||
response.sendError(HttpServletResponse.SC_INTERNAL_SERVER_ERROR,
|
||||
"Error processing response: " + e.getMessage());
|
||||
}
|
||||
catch (IOException ex) {
|
||||
logger.error("Failed to send error response: {}", ex.getMessage());
|
||||
}
|
||||
}
|
||||
}, error -> {
|
||||
try {
|
||||
logger.error("Error processing message: {}", error.getMessage());
|
||||
McpError mcpError = new McpError(error.getMessage());
|
||||
response.setContentType("application/json");
|
||||
response.setCharacterEncoding("UTF-8");
|
||||
response.setStatus(HttpServletResponse.SC_INTERNAL_SERVER_ERROR);
|
||||
String jsonError = objectMapper.writeValueAsString(mcpError);
|
||||
PrintWriter writer = response.getWriter();
|
||||
writer.write(jsonError);
|
||||
writer.flush();
|
||||
}
|
||||
catch (IOException e) {
|
||||
logger.error("Failed to send error response: {}", e.getMessage());
|
||||
try {
|
||||
response.sendError(HttpServletResponse.SC_INTERNAL_SERVER_ERROR,
|
||||
"Error sending error response: " + e.getMessage());
|
||||
}
|
||||
catch (IOException ex) {
|
||||
logger.error("Failed to send error response: {}", ex.getMessage());
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
else {
|
||||
response.sendError(HttpServletResponse.SC_SERVICE_UNAVAILABLE, "No message handler configured");
|
||||
}
|
||||
}
|
||||
catch (Exception e) {
|
||||
logger.error("Invalid message format: {}", e.getMessage());
|
||||
try {
|
||||
McpError mcpError = new McpError("Invalid message format: " + e.getMessage());
|
||||
response.setContentType("application/json");
|
||||
response.setCharacterEncoding("UTF-8");
|
||||
response.setStatus(HttpServletResponse.SC_BAD_REQUEST);
|
||||
String jsonError = objectMapper.writeValueAsString(mcpError);
|
||||
PrintWriter writer = response.getWriter();
|
||||
writer.write(jsonError);
|
||||
writer.flush();
|
||||
}
|
||||
catch (IOException ex) {
|
||||
logger.error("Failed to send error response: {}", ex.getMessage());
|
||||
response.sendError(HttpServletResponse.SC_BAD_REQUEST, "Invalid message format");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public Mono<Void> connect(Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> handler) {
|
||||
this.connectHandler = handler;
|
||||
return Mono.empty();
|
||||
}
|
||||
|
||||
@Override
|
||||
public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message) {
|
||||
if (sessions.isEmpty()) {
|
||||
logger.debug("No active sessions to broadcast message to");
|
||||
return Mono.empty();
|
||||
}
|
||||
|
||||
return Mono.create(sink -> {
|
||||
try {
|
||||
String jsonText = objectMapper.writeValueAsString(message);
|
||||
|
||||
sessions.values().forEach(session -> {
|
||||
try {
|
||||
sendEvent(session.writer, MESSAGE_EVENT_TYPE, jsonText);
|
||||
}
|
||||
catch (IOException e) {
|
||||
logger.error("Failed to send message to session {}: {}", session.id, e.getMessage());
|
||||
removeSession(session);
|
||||
}
|
||||
});
|
||||
|
||||
sink.success();
|
||||
}
|
||||
catch (Exception e) {
|
||||
logger.error("Failed to process message: {}", e.getMessage());
|
||||
sink.error(new McpError("Failed to process message: " + e.getMessage()));
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
@Override
|
||||
public void close() {
|
||||
ServerMcpTransport.super.close();
|
||||
}
|
||||
|
||||
@Override
|
||||
public <T> T unmarshalFrom(Object data, TypeReference<T> typeRef) {
|
||||
return objectMapper.convertValue(data, typeRef);
|
||||
}
|
||||
|
||||
@Override
|
||||
public Mono<Void> closeGracefully() {
|
||||
isClosing.set(true);
|
||||
logger.debug("Initiating graceful shutdown with {} active sessions", sessions.size());
|
||||
|
||||
return Mono.create(sink -> {
|
||||
sessions.values().forEach(this::removeSession);
|
||||
sink.success();
|
||||
});
|
||||
}
|
||||
|
||||
private void sendEvent(PrintWriter writer, String eventType, String data) throws IOException {
|
||||
writer.write("event: " + eventType + "\n");
|
||||
writer.write("data: " + data + "\n\n");
|
||||
writer.flush();
|
||||
|
||||
if (writer.checkError()) {
|
||||
throw new IOException("Client disconnected");
|
||||
}
|
||||
}
|
||||
|
||||
private void removeSession(ClientSession session) {
|
||||
sessions.remove(session.id);
|
||||
session.asyncContext.complete();
|
||||
}
|
||||
|
||||
private static class ClientSession {
|
||||
|
||||
private final String id;
|
||||
|
||||
private final AsyncContext asyncContext;
|
||||
|
||||
private final PrintWriter writer;
|
||||
|
||||
ClientSession(String id, AsyncContext asyncContext, PrintWriter writer) {
|
||||
this.id = id;
|
||||
this.asyncContext = asyncContext;
|
||||
this.writer = writer;
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
@Override
|
||||
public void destroy() {
|
||||
closeGracefully().block();
|
||||
super.destroy();
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,81 @@
|
|||
/*
|
||||
* 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.attic;
|
||||
|
||||
import java.time.Duration;
|
||||
import java.util.Map;
|
||||
|
||||
import org.springframework.ai.mcp.client.McpClient;
|
||||
import org.springframework.ai.mcp.client.McpSyncClient;
|
||||
import org.springframework.ai.mcp.client.transport.SseClientTransport;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.CallToolRequest;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.CallToolResult;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.ListPromptsResult;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.ListResourcesResult;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.ListToolsResult;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.Resource;
|
||||
import org.springframework.web.reactive.function.client.WebClient;
|
||||
|
||||
/**
|
||||
* @author Christian Tzolov
|
||||
* @since 1.0.0
|
||||
*/
|
||||
|
||||
public class MainSSE {
|
||||
|
||||
public static void main(String[] args) {
|
||||
|
||||
try (McpSyncClient mcpClient = McpClient
|
||||
.using(new SseClientTransport(WebClient.builder().baseUrl("http://localhost:3001")))
|
||||
.requestTimeout(Duration.ofSeconds(1000))
|
||||
.sync()) {
|
||||
|
||||
mcpClient.initialize();
|
||||
|
||||
ListToolsResult tools = mcpClient.listTools();
|
||||
System.out.println("Tools: " + tools);
|
||||
|
||||
mcpClient.ping();
|
||||
|
||||
CallToolRequest callToolRequest = new CallToolRequest("echo", Map.of("message", "Hello MCP Spring AI!"));
|
||||
CallToolResult callToolResult = mcpClient.callTool(callToolRequest);
|
||||
System.out.println("Call Tool Result: " + callToolResult);
|
||||
|
||||
// mcpClient.sendRootsListChanged();
|
||||
|
||||
// Resources
|
||||
ListResourcesResult resources = mcpClient.listResources();
|
||||
System.out.println("Resources Size: " + resources.resources().size());
|
||||
System.out.println("Resources: " + resources);
|
||||
for (Resource resource : resources.resources()) {
|
||||
System.out.println(mcpClient.readResource(resource));
|
||||
}
|
||||
|
||||
var resourceTemplate = mcpClient.listResourceTemplates();
|
||||
System.out.println("Resource Templates: " + resourceTemplate);
|
||||
|
||||
ListPromptsResult prompts = mcpClient.listPrompts();
|
||||
for (var prompt : prompts.prompts()) {
|
||||
System.out.println("Prompt: " + prompt);
|
||||
}
|
||||
|
||||
}
|
||||
catch (Exception e) {
|
||||
e.printStackTrace();
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -1,61 +0,0 @@
|
|||
/*
|
||||
* 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.client;
|
||||
|
||||
import org.junit.jupiter.api.Timeout;
|
||||
import org.testcontainers.containers.GenericContainer;
|
||||
import org.testcontainers.containers.wait.strategy.Wait;
|
||||
|
||||
import org.springframework.ai.mcp.client.transport.SseClientTransport;
|
||||
import org.springframework.ai.mcp.spec.ClientMcpTransport;
|
||||
import org.springframework.web.reactive.function.client.WebClient;
|
||||
|
||||
/**
|
||||
* Tests for the {@link McpAsyncClient} with {@link SseClientTransport}.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
*/
|
||||
@Timeout(15) // Giving extra time beyond the client timeout
|
||||
class SseMcpAsyncClientTests extends AbstractMcpAsyncClientTests {
|
||||
|
||||
static String host = "http://localhost:3001";
|
||||
|
||||
// Uses the https://github.com/tzolov/mcp-everything-server-docker-image
|
||||
@SuppressWarnings("resource")
|
||||
GenericContainer<?> container = new GenericContainer<>("docker.io/tzolov/mcp-everything-server:v1")
|
||||
.withLogConsumer(outputFrame -> System.out.println(outputFrame.getUtf8String()))
|
||||
.withExposedPorts(3001)
|
||||
.waitingFor(Wait.forHttp("/").forStatusCode(404));
|
||||
|
||||
@Override
|
||||
protected ClientMcpTransport createMcpTransport() {
|
||||
return new SseClientTransport(WebClient.builder().baseUrl(host));
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void onStart() {
|
||||
container.start();
|
||||
int port = container.getMappedPort(3001);
|
||||
host = "http://" + container.getHost() + ":" + port;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void onClose() {
|
||||
container.stop();
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -1,61 +0,0 @@
|
|||
/*
|
||||
* 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.client;
|
||||
|
||||
import org.junit.jupiter.api.Timeout;
|
||||
import org.testcontainers.containers.GenericContainer;
|
||||
import org.testcontainers.containers.wait.strategy.Wait;
|
||||
|
||||
import org.springframework.ai.mcp.client.transport.SseClientTransport;
|
||||
import org.springframework.ai.mcp.spec.ClientMcpTransport;
|
||||
import org.springframework.web.reactive.function.client.WebClient;
|
||||
|
||||
/**
|
||||
* Tests for the {@link McpSyncClient} with {@link SseClientTransport}.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
*/
|
||||
@Timeout(15) // Giving extra time beyond the client timeout
|
||||
class SseMcpSyncClientTests extends AbstractMcpSyncClientTests {
|
||||
|
||||
static String host = "http://localhost:3001";
|
||||
|
||||
// Uses the https://github.com/tzolov/mcp-everything-server-docker-image
|
||||
@SuppressWarnings("resource")
|
||||
GenericContainer<?> container = new GenericContainer<>("docker.io/tzolov/mcp-everything-server:v1")
|
||||
.withLogConsumer(outputFrame -> System.out.println(outputFrame.getUtf8String()))
|
||||
.withExposedPorts(3001)
|
||||
.waitingFor(Wait.forHttp("/").forStatusCode(404));
|
||||
|
||||
@Override
|
||||
protected ClientMcpTransport createMcpTransport() {
|
||||
return new SseClientTransport(WebClient.builder().baseUrl(host));
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void onStart() {
|
||||
container.start();
|
||||
int port = container.getMappedPort(3001);
|
||||
host = "http://" + container.getHost() + ":" + port;
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void onClose() {
|
||||
container.stop();
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,290 @@
|
|||
/*
|
||||
* 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.client.transport;
|
||||
|
||||
import java.time.Duration;
|
||||
import java.util.Map;
|
||||
import java.util.concurrent.atomic.AtomicInteger;
|
||||
import java.util.function.Function;
|
||||
|
||||
import org.junit.jupiter.api.AfterEach;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.Timeout;
|
||||
import org.testcontainers.containers.GenericContainer;
|
||||
import org.testcontainers.containers.wait.strategy.Wait;
|
||||
import reactor.core.publisher.Mono;
|
||||
import reactor.core.publisher.Sinks;
|
||||
import reactor.test.StepVerifier;
|
||||
|
||||
import org.springframework.ai.mcp.spec.McpSchema;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.JSONRPCRequest;
|
||||
import org.springframework.http.codec.ServerSentEvent;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.assertj.core.api.Assertions.assertThatCode;
|
||||
|
||||
/**
|
||||
* Tests for the {@link HttpClientSseClientTransport} class.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
*/
|
||||
@Timeout(15)
|
||||
class HttpClientSseClientTransportTests {
|
||||
|
||||
static String host = "http://localhost:3001";
|
||||
|
||||
@SuppressWarnings("resource")
|
||||
GenericContainer<?> container = new GenericContainer<>("docker.io/tzolov/mcp-everything-server:v1")
|
||||
.withLogConsumer(outputFrame -> System.out.println(outputFrame.getUtf8String()))
|
||||
.withExposedPorts(3001)
|
||||
.waitingFor(Wait.forHttp("/").forStatusCode(404));
|
||||
|
||||
private TestHttpClientSseClientTransport transport;
|
||||
|
||||
// Test class to access protected methods
|
||||
static class TestHttpClientSseClientTransport extends HttpClientSseClientTransport {
|
||||
|
||||
private final AtomicInteger inboundMessageCount = new AtomicInteger(0);
|
||||
|
||||
private Sinks.Many<ServerSentEvent<String>> events = Sinks.many().unicast().onBackpressureBuffer();
|
||||
|
||||
public TestHttpClientSseClientTransport(String baseUri) {
|
||||
super(baseUri);
|
||||
}
|
||||
|
||||
public int getInboundMessageCount() {
|
||||
return inboundMessageCount.get();
|
||||
}
|
||||
|
||||
public void simulateEndpointEvent(String jsonMessage) {
|
||||
events.tryEmitNext(ServerSentEvent.<String>builder().event("endpoint").data(jsonMessage).build());
|
||||
inboundMessageCount.incrementAndGet();
|
||||
}
|
||||
|
||||
public void simulateMessageEvent(String jsonMessage) {
|
||||
events.tryEmitNext(ServerSentEvent.<String>builder().event("message").data(jsonMessage).build());
|
||||
inboundMessageCount.incrementAndGet();
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
void startContainer() {
|
||||
container.start();
|
||||
int port = container.getMappedPort(3001);
|
||||
host = "http://" + container.getHost() + ":" + port;
|
||||
}
|
||||
|
||||
@BeforeEach
|
||||
void setUp() {
|
||||
startContainer();
|
||||
transport = new TestHttpClientSseClientTransport(host);
|
||||
transport.connect(Function.identity()).block();
|
||||
}
|
||||
|
||||
@AfterEach
|
||||
void afterEach() {
|
||||
if (transport != null) {
|
||||
assertThatCode(() -> transport.closeGracefully().block(Duration.ofSeconds(10))).doesNotThrowAnyException();
|
||||
}
|
||||
cleanup();
|
||||
}
|
||||
|
||||
void cleanup() {
|
||||
container.stop();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testMessageProcessing() {
|
||||
// Create a test message
|
||||
JSONRPCRequest testMessage = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "test-method", "test-id",
|
||||
Map.of("key", "value"));
|
||||
|
||||
// Simulate receiving the message
|
||||
transport.simulateMessageEvent("""
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"method": "test-method",
|
||||
"id": "test-id",
|
||||
"params": {"key": "value"}
|
||||
}
|
||||
""");
|
||||
|
||||
// Subscribe to messages and verify
|
||||
StepVerifier.create(transport.sendMessage(testMessage)).verifyComplete();
|
||||
|
||||
assertThat(transport.getInboundMessageCount()).isEqualTo(1);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testResponseMessageProcessing() {
|
||||
// Simulate receiving a response message
|
||||
transport.simulateMessageEvent("""
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": "test-id",
|
||||
"result": {"status": "success"}
|
||||
}
|
||||
""");
|
||||
|
||||
// Create and send a request message
|
||||
JSONRPCRequest testMessage = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "test-method", "test-id",
|
||||
Map.of("key", "value"));
|
||||
|
||||
// Verify message handling
|
||||
StepVerifier.create(transport.sendMessage(testMessage)).verifyComplete();
|
||||
|
||||
assertThat(transport.getInboundMessageCount()).isEqualTo(1);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testErrorMessageProcessing() {
|
||||
// Simulate receiving an error message
|
||||
transport.simulateMessageEvent("""
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": "test-id",
|
||||
"error": {
|
||||
"code": -32600,
|
||||
"message": "Invalid Request"
|
||||
}
|
||||
}
|
||||
""");
|
||||
|
||||
// Create and send a request message
|
||||
JSONRPCRequest testMessage = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "test-method", "test-id",
|
||||
Map.of("key", "value"));
|
||||
|
||||
// Verify message handling
|
||||
StepVerifier.create(transport.sendMessage(testMessage)).verifyComplete();
|
||||
|
||||
assertThat(transport.getInboundMessageCount()).isEqualTo(1);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testNotificationMessageProcessing() {
|
||||
// Simulate receiving a notification message (no id)
|
||||
transport.simulateMessageEvent("""
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"method": "update",
|
||||
"params": {"status": "processing"}
|
||||
}
|
||||
""");
|
||||
|
||||
// Verify the notification was processed
|
||||
assertThat(transport.getInboundMessageCount()).isEqualTo(1);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testGracefulShutdown() {
|
||||
// Test graceful shutdown
|
||||
StepVerifier.create(transport.closeGracefully()).verifyComplete();
|
||||
|
||||
// Create a test message
|
||||
JSONRPCRequest testMessage = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "test-method", "test-id",
|
||||
Map.of("key", "value"));
|
||||
|
||||
// Verify message is not processed after shutdown
|
||||
StepVerifier.create(transport.sendMessage(testMessage)).verifyComplete();
|
||||
|
||||
// Message count should remain 0 after shutdown
|
||||
assertThat(transport.getInboundMessageCount()).isEqualTo(0);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testRetryBehavior() {
|
||||
// Create a client that simulates connection failures
|
||||
HttpClientSseClientTransport failingTransport = new HttpClientSseClientTransport("http://non-existent-host");
|
||||
|
||||
// Verify that the transport attempts to reconnect
|
||||
StepVerifier.create(Mono.delay(Duration.ofSeconds(2))).expectNextCount(1).verifyComplete();
|
||||
|
||||
// Clean up
|
||||
failingTransport.closeGracefully().block();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testMultipleMessageProcessing() {
|
||||
// Simulate receiving multiple messages in sequence
|
||||
transport.simulateMessageEvent("""
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"method": "method1",
|
||||
"id": "id1",
|
||||
"params": {"key": "value1"}
|
||||
}
|
||||
""");
|
||||
|
||||
transport.simulateMessageEvent("""
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"method": "method2",
|
||||
"id": "id2",
|
||||
"params": {"key": "value2"}
|
||||
}
|
||||
""");
|
||||
|
||||
// Create and send corresponding messages
|
||||
JSONRPCRequest message1 = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "method1", "id1",
|
||||
Map.of("key", "value1"));
|
||||
|
||||
JSONRPCRequest message2 = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "method2", "id2",
|
||||
Map.of("key", "value2"));
|
||||
|
||||
// Verify both messages are processed
|
||||
StepVerifier.create(transport.sendMessage(message1).then(transport.sendMessage(message2))).verifyComplete();
|
||||
|
||||
// Verify message count
|
||||
assertThat(transport.getInboundMessageCount()).isEqualTo(2);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testMessageOrderPreservation() {
|
||||
// Simulate receiving messages in a specific order
|
||||
transport.simulateMessageEvent("""
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"method": "first",
|
||||
"id": "1",
|
||||
"params": {"sequence": 1}
|
||||
}
|
||||
""");
|
||||
|
||||
transport.simulateMessageEvent("""
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"method": "second",
|
||||
"id": "2",
|
||||
"params": {"sequence": 2}
|
||||
}
|
||||
""");
|
||||
|
||||
transport.simulateMessageEvent("""
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"method": "third",
|
||||
"id": "3",
|
||||
"params": {"sequence": 3}
|
||||
}
|
||||
""");
|
||||
|
||||
// Verify message count and order
|
||||
assertThat(transport.getInboundMessageCount()).isEqualTo(3);
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -18,7 +18,6 @@ package org.springframework.ai.mcp.server;
|
|||
|
||||
import java.time.Duration;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
import org.junit.jupiter.api.AfterEach;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
|
|
|
|||
|
|
@ -1,66 +0,0 @@
|
|||
/*
|
||||
* 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 com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import org.junit.jupiter.api.Timeout;
|
||||
import reactor.netty.DisposableServer;
|
||||
import reactor.netty.http.server.HttpServer;
|
||||
|
||||
import org.springframework.ai.mcp.server.transport.SseServerTransport;
|
||||
import org.springframework.ai.mcp.spec.ServerMcpTransport;
|
||||
import org.springframework.http.server.reactive.HttpHandler;
|
||||
import org.springframework.http.server.reactive.ReactorHttpHandlerAdapter;
|
||||
import org.springframework.web.reactive.function.server.RouterFunctions;
|
||||
|
||||
/**
|
||||
* Tests for {@link McpAsyncServer} using {@link SseServerTransport}.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
*/
|
||||
@Timeout(15) // Giving extra time beyond the client timeout
|
||||
class SseMcpAsyncServerTests extends AbstractMcpAsyncServerTests {
|
||||
|
||||
private static final int PORT = 8181;
|
||||
|
||||
private static final String MESSAGE_ENDPOINT = "/mcp/message";
|
||||
|
||||
private DisposableServer httpServer;
|
||||
|
||||
@Override
|
||||
protected ServerMcpTransport createMcpTransport() {
|
||||
var transport = new SseServerTransport(new ObjectMapper(), MESSAGE_ENDPOINT);
|
||||
|
||||
HttpHandler httpHandler = RouterFunctions.toHttpHandler(transport.getRouterFunction());
|
||||
ReactorHttpHandlerAdapter adapter = new ReactorHttpHandlerAdapter(httpHandler);
|
||||
httpServer = HttpServer.create().port(PORT).handle(adapter).bindNow();
|
||||
|
||||
return transport;
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void onStart() {
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void onClose() {
|
||||
if (httpServer != null) {
|
||||
httpServer.disposeNow();
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -1,66 +0,0 @@
|
|||
/*
|
||||
* 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 com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import org.junit.jupiter.api.Timeout;
|
||||
import reactor.netty.DisposableServer;
|
||||
import reactor.netty.http.server.HttpServer;
|
||||
|
||||
import org.springframework.ai.mcp.server.transport.SseServerTransport;
|
||||
import org.springframework.ai.mcp.spec.ServerMcpTransport;
|
||||
import org.springframework.http.server.reactive.HttpHandler;
|
||||
import org.springframework.http.server.reactive.ReactorHttpHandlerAdapter;
|
||||
import org.springframework.web.reactive.function.server.RouterFunctions;
|
||||
|
||||
/**
|
||||
* Tests for {@link McpSyncServer} using {@link SseServerTransport}.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
*/
|
||||
@Timeout(15) // Giving extra time beyond the client timeout
|
||||
class SseMcpSyncServerTests extends AbstractMcpSyncServerTests {
|
||||
|
||||
private static final int PORT = 8182;
|
||||
|
||||
private static final String MESSAGE_ENDPOINT = "/mcp/message";
|
||||
|
||||
private DisposableServer httpServer;
|
||||
|
||||
private SseServerTransport transport;
|
||||
|
||||
@Override
|
||||
protected ServerMcpTransport createMcpTransport() {
|
||||
transport = new SseServerTransport(new ObjectMapper(), MESSAGE_ENDPOINT);
|
||||
return transport;
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void onStart() {
|
||||
HttpHandler httpHandler = RouterFunctions.toHttpHandler(transport.getRouterFunction());
|
||||
ReactorHttpHandlerAdapter adapter = new ReactorHttpHandlerAdapter(httpHandler);
|
||||
httpServer = HttpServer.create().port(PORT).handle(adapter).bindNow();
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void onClose() {
|
||||
if (httpServer != null) {
|
||||
httpServer.disposeNow();
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,439 @@
|
|||
/*
|
||||
* 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.junit.jupiter.api.AfterEach;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
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.SseClientTransport;
|
||||
import org.springframework.ai.mcp.server.McpServer.ToolRegistration;
|
||||
import org.springframework.ai.mcp.server.transport.SseServerTransport;
|
||||
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.http.server.reactive.HttpHandler;
|
||||
import org.springframework.http.server.reactive.ReactorHttpHandlerAdapter;
|
||||
import org.springframework.web.client.RestClient;
|
||||
import org.springframework.web.reactive.function.client.WebClient;
|
||||
import org.springframework.web.reactive.function.server.RouterFunctions;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.assertj.core.api.Assertions.assertThatThrownBy;
|
||||
import static org.awaitility.Awaitility.await;
|
||||
|
||||
public class SseAsyncIntegrationTests {
|
||||
|
||||
private static final int PORT = 8181;
|
||||
|
||||
private static final String MESSAGE_ENDPOINT = "/mcp/message";
|
||||
|
||||
private DisposableServer httpServer;
|
||||
|
||||
private SseServerTransport mcpServerTransport;
|
||||
|
||||
McpClient.Builder clientBuilder;
|
||||
|
||||
@BeforeEach
|
||||
public void before() {
|
||||
this.mcpServerTransport = new SseServerTransport(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 SseClientTransport(WebClient.builder().baseUrl("http://localhost:" + PORT)));
|
||||
}
|
||||
|
||||
@AfterEach
|
||||
public void after() {
|
||||
if (httpServer != null) {
|
||||
httpServer.disposeNow();
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------
|
||||
// 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<CreateMessageRequest, CreateMessageResult> 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<Root> roots = List.of(new Root("uri1://", "root1"), new Root("uri2://", "root2"));
|
||||
|
||||
AtomicReference<List<Root>> 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
|
||||
var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().build()) // No
|
||||
// roots
|
||||
// capability
|
||||
.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<List<Root>> 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<Root> roots = List.of(new Root("uri1://", "root1"));
|
||||
|
||||
AtomicReference<List<Root>> rootsRef1 = new AtomicReference<>();
|
||||
AtomicReference<List<Root>> 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<Root> roots = List.of(new Root("uri1://", "root1"));
|
||||
|
||||
AtomicReference<List<Root>> 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<List<Tool>> 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();
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -1,283 +0,0 @@
|
|||
package org.springframework.ai.mcp.server.transport;
|
||||
|
||||
import java.time.Duration;
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.concurrent.CountDownLatch;
|
||||
import java.util.concurrent.TimeUnit;
|
||||
import java.util.concurrent.atomic.AtomicReference;
|
||||
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Disabled;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.Timeout;
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
import reactor.core.Disposable;
|
||||
import reactor.core.publisher.Flux;
|
||||
import reactor.core.publisher.Mono;
|
||||
import reactor.test.StepVerifier;
|
||||
|
||||
import org.springframework.ai.mcp.spec.McpSchema;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.JSONRPCRequest;
|
||||
import org.springframework.http.HttpStatus;
|
||||
import org.springframework.http.MediaType;
|
||||
import org.springframework.http.codec.ServerSentEvent;
|
||||
import org.springframework.test.web.reactive.server.WebTestClient;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.assertj.core.api.Assertions.assertThatThrownBy;
|
||||
|
||||
/**
|
||||
* Tests for the {@link SseServerTransport} class.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
*/
|
||||
@Timeout(15)
|
||||
class SseServerTransportTests {
|
||||
|
||||
private static final Logger logger = LoggerFactory.getLogger(SseServerTransportTests.class);
|
||||
|
||||
private ObjectMapper objectMapper;
|
||||
|
||||
private String messageEndpoint;
|
||||
|
||||
private SseServerTransport transport;
|
||||
|
||||
private WebTestClient webTestClient;
|
||||
|
||||
@BeforeEach
|
||||
void setUp() {
|
||||
objectMapper = new ObjectMapper();
|
||||
messageEndpoint = "/message";
|
||||
transport = new SseServerTransport(objectMapper, messageEndpoint);
|
||||
webTestClient = WebTestClient.bindToRouterFunction(transport.getRouterFunction()).build();
|
||||
}
|
||||
|
||||
@Test
|
||||
void constructorValidation() {
|
||||
assertThatThrownBy(() -> new SseServerTransport(null, "/message")).isInstanceOf(IllegalArgumentException.class)
|
||||
.hasMessageContaining("ObjectMapper must not be null");
|
||||
|
||||
assertThatThrownBy(() -> new SseServerTransport(new ObjectMapper(), null))
|
||||
.isInstanceOf(IllegalArgumentException.class)
|
||||
.hasMessageContaining("Message endpoint must not be null");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testSseConnectionEstablishment() {
|
||||
List<ServerSentEvent<String>> events = new ArrayList<>();
|
||||
|
||||
webTestClient.get()
|
||||
.uri("/sse")
|
||||
.accept(MediaType.TEXT_EVENT_STREAM)
|
||||
.exchange()
|
||||
.expectStatus()
|
||||
.isOk()
|
||||
.expectHeader()
|
||||
.contentTypeCompatibleWith(MediaType.TEXT_EVENT_STREAM)
|
||||
.returnResult(String.class)
|
||||
.getResponseBody()
|
||||
.map(data -> ServerSentEvent.<String>builder().data(data).build())
|
||||
.take(1) // Take only the initial endpoint event
|
||||
.subscribe(events::add);
|
||||
|
||||
// Wait a bit for the event to be received
|
||||
StepVerifier.create(Mono.delay(Duration.ofMillis(500))).expectNextCount(1).verifyComplete();
|
||||
|
||||
assertThat(events).hasSize(1);
|
||||
assertThat(events.get(0).data()).isEqualTo(messageEndpoint);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testMessageHandling() {
|
||||
AtomicReference<McpSchema.JSONRPCMessage> receivedMessage = new AtomicReference<>();
|
||||
|
||||
// Set up message handler
|
||||
transport.connect(message -> {
|
||||
message.doOnNext(receivedMessage::set).subscribe();
|
||||
return Mono.empty();
|
||||
}).block();
|
||||
|
||||
// Create a test message
|
||||
JSONRPCRequest testMessage = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "test-method", "test-id",
|
||||
Map.of("key", "value"));
|
||||
|
||||
// Send message to endpoint
|
||||
webTestClient.post()
|
||||
.uri(messageEndpoint)
|
||||
.contentType(MediaType.APPLICATION_JSON)
|
||||
.bodyValue(testMessage)
|
||||
.exchange()
|
||||
.expectStatus()
|
||||
.isOk();
|
||||
|
||||
// Verify the message was received and processed
|
||||
assertThat(receivedMessage.get()).isNotNull();
|
||||
McpSchema.JSONRPCRequest receivedRequest = (McpSchema.JSONRPCRequest) receivedMessage.get();
|
||||
assertThat(receivedRequest.id()).isEqualTo(testMessage.id());
|
||||
assertThat(receivedRequest.method()).isEqualTo(testMessage.method());
|
||||
}
|
||||
|
||||
@Test
|
||||
@Disabled("Flaky test")
|
||||
void testBroadcastMessage() {
|
||||
// Create test clients
|
||||
int clientCount = 3;
|
||||
CountDownLatch connectLatch = new CountDownLatch(clientCount);
|
||||
CountDownLatch messageLatch = new CountDownLatch(clientCount);
|
||||
|
||||
List<List<ServerSentEvent<String>>> allReceivedEvents = new ArrayList<>();
|
||||
List<Disposable> subscriptions = new ArrayList<>();
|
||||
|
||||
// Connect clients
|
||||
for (int i = 0; i < clientCount; i++) {
|
||||
List<ServerSentEvent<String>> clientEvents = new ArrayList<>();
|
||||
allReceivedEvents.add(clientEvents);
|
||||
|
||||
Flux<ServerSentEvent<String>> eventStream = webTestClient.get()
|
||||
.uri("/sse")
|
||||
.accept(MediaType.TEXT_EVENT_STREAM)
|
||||
.exchange()
|
||||
.expectStatus()
|
||||
.isOk()
|
||||
.returnResult(String.class)
|
||||
.getResponseBody()
|
||||
.map(data -> ServerSentEvent.<String>builder().data(data).build());
|
||||
|
||||
Disposable subscription = eventStream.doOnNext(event -> {
|
||||
clientEvents.add(event);
|
||||
if (event.event() != null && event.event().equals(SseServerTransport.ENDPOINT_EVENT_TYPE)) {
|
||||
connectLatch.countDown();
|
||||
}
|
||||
else if (event.event() != null && event.event().equals(SseServerTransport.MESSAGE_EVENT_TYPE)) {
|
||||
messageLatch.countDown();
|
||||
}
|
||||
}).subscribe();
|
||||
|
||||
subscriptions.add(subscription);
|
||||
}
|
||||
|
||||
// Wait for all clients to connect
|
||||
try {
|
||||
assertThat(connectLatch.await(5, TimeUnit.SECONDS)).isTrue();
|
||||
}
|
||||
catch (InterruptedException e) {
|
||||
throw new RuntimeException(e);
|
||||
}
|
||||
|
||||
// Verify initial connections
|
||||
for (List<ServerSentEvent<String>> events : allReceivedEvents) {
|
||||
assertThat(events).hasSize(1);
|
||||
assertThat(events.get(0).data()).isEqualTo(messageEndpoint);
|
||||
}
|
||||
|
||||
// Give clients time to fully establish their subscriptions
|
||||
logger.debug("Waiting for subscriptions to stabilize...");
|
||||
try {
|
||||
Thread.sleep(1000);
|
||||
}
|
||||
catch (InterruptedException e) {
|
||||
throw new RuntimeException(e);
|
||||
}
|
||||
logger.debug("Sending broadcast message to {} clients", clientCount);
|
||||
|
||||
// Send broadcast message
|
||||
JSONRPCRequest broadcastMessage = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "broadcast", "broadcast-id",
|
||||
Map.of("message", "Hello all!"));
|
||||
|
||||
// Send the message
|
||||
transport.sendMessage(broadcastMessage).block(Duration.ofSeconds(5));
|
||||
|
||||
// Wait for all clients to receive the broadcast
|
||||
try {
|
||||
assertThat(messageLatch.await(5, TimeUnit.SECONDS)).isTrue();
|
||||
}
|
||||
catch (InterruptedException e) {
|
||||
throw new RuntimeException(e);
|
||||
}
|
||||
|
||||
// Verify each client received both messages
|
||||
for (List<ServerSentEvent<String>> events : allReceivedEvents) {
|
||||
assertThat(events).hasSize(2);
|
||||
assertThat(events.get(0).data()).isEqualTo(messageEndpoint);
|
||||
assertThat(events.get(1).data()).contains("broadcast-id");
|
||||
}
|
||||
|
||||
// Cleanup
|
||||
subscriptions.forEach(Disposable::dispose);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testGracefulShutdown() {
|
||||
// Connect a client
|
||||
Flux<ServerSentEvent<String>> eventStream = webTestClient.get()
|
||||
.uri("/sse")
|
||||
.accept(MediaType.TEXT_EVENT_STREAM)
|
||||
.exchange()
|
||||
.expectStatus()
|
||||
.isOk()
|
||||
.returnResult(String.class)
|
||||
.getResponseBody()
|
||||
.map(data -> ServerSentEvent.<String>builder().data(data).build());
|
||||
|
||||
List<ServerSentEvent<String>> receivedEvents = new ArrayList<>();
|
||||
eventStream.subscribe(receivedEvents::add);
|
||||
|
||||
// Wait for connection
|
||||
StepVerifier.create(Mono.delay(Duration.ofMillis(500))).expectNextCount(1).verifyComplete();
|
||||
|
||||
// Initiate shutdown
|
||||
transport.closeGracefully().block(Duration.ofSeconds(5));
|
||||
|
||||
// Verify server rejects new connections with timeout
|
||||
webTestClient.get()
|
||||
.uri("/sse")
|
||||
.accept(MediaType.TEXT_EVENT_STREAM)
|
||||
.exchange()
|
||||
.expectStatus()
|
||||
.isEqualTo(HttpStatus.SERVICE_UNAVAILABLE)
|
||||
.expectBody(String.class)
|
||||
.isEqualTo("Server is shutting down");
|
||||
|
||||
// Verify server rejects new messages with timeout
|
||||
webTestClient.post()
|
||||
.uri(messageEndpoint)
|
||||
.contentType(MediaType.APPLICATION_JSON)
|
||||
.bodyValue("""
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"method": "test",
|
||||
"id": "1",
|
||||
"params": {}
|
||||
}
|
||||
""")
|
||||
.exchange()
|
||||
.expectStatus()
|
||||
.isEqualTo(HttpStatus.SERVICE_UNAVAILABLE)
|
||||
.expectBody(String.class)
|
||||
.isEqualTo("Server is shutting down");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testInvalidMessageHandling() {
|
||||
// Test invalid JSON
|
||||
webTestClient.post()
|
||||
.uri(messageEndpoint)
|
||||
.contentType(MediaType.APPLICATION_JSON)
|
||||
.bodyValue("invalid json")
|
||||
.exchange()
|
||||
.expectStatus()
|
||||
.isBadRequest();
|
||||
|
||||
// Test invalid message format
|
||||
webTestClient.post().uri(messageEndpoint).contentType(MediaType.APPLICATION_JSON).bodyValue("""
|
||||
{
|
||||
"invalid": "message"
|
||||
}
|
||||
""").exchange().expectStatus().isBadRequest();
|
||||
}
|
||||
|
||||
}
|
||||
Loading…
Reference in a new issue