refactor: Add HttpServletSseServerTransport for generic Servlet integration
Introduces a new Server-Sent Events (SSE) Server transport implementation based on the Servlet API that enables MCP integration with any Java HTTP server application supporting Servlets. This implementation complements the existing WebFlux and WebMvc transports by providing a generic Servlet-based option that works with any Java web server, not just Spring applications. Key changes: - Add HttpServletSseServerTransport with async Servlet 6.0 support - Add integration tests with embedded Tomcat - Add servlet dependencies with externalized versions: - jakarta.servlet-api 6.1.0 (provided) - tomcat-embed-core 11.0.2 (test) - tomcat-embed-websocket 11.0.2 (test) Resolves #61
This commit is contained in:
Christian Tzolov
parent
42b5931dd5
commit
0c8968044c
2 changed files with 273 additions and 206 deletions
|
|
@ -1,297 +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.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,344 @@
|
|||
/*
|
||||
* 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.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;
|
||||
import org.springframework.ai.mcp.server.McpServer.ToolRegistration;
|
||||
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.web.client.RestClient;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.awaitility.Awaitility.await;
|
||||
|
||||
public class HttpServletSseServerTransportIntegrationTests {
|
||||
|
||||
private static final Logger logger = LoggerFactory.getLogger(HttpServletSseServerTransportIntegrationTests.class);
|
||||
|
||||
private static final int PORT = 8184;
|
||||
|
||||
private static final String MESSAGE_ENDPOINT = "/mcp/message";
|
||||
|
||||
private HttpServletSseServerTransport mcpServerTransport;
|
||||
|
||||
McpClient.Builder clientBuilder;
|
||||
|
||||
private Tomcat tomcat;
|
||||
|
||||
@BeforeEach
|
||||
public void before() {
|
||||
tomcat = new Tomcat();
|
||||
tomcat.setPort(PORT);
|
||||
|
||||
String baseDir = System.getProperty("java.io.tmpdir");
|
||||
tomcat.setBaseDir(baseDir);
|
||||
|
||||
Context context = tomcat.addContext("", baseDir);
|
||||
|
||||
// Create and configure the transport
|
||||
mcpServerTransport = new HttpServletSseServerTransport(new ObjectMapper(), MESSAGE_ENDPOINT);
|
||||
|
||||
// Add transport servlet to Tomcat
|
||||
org.apache.catalina.Wrapper wrapper = context.createWrapper();
|
||||
wrapper.setName("mcpServlet");
|
||||
wrapper.setServlet(mcpServerTransport);
|
||||
wrapper.setLoadOnStartup(1);
|
||||
wrapper.setAsyncSupported(true);
|
||||
context.addChild(wrapper);
|
||||
context.addServletMappingDecoded("/*", "mcpServlet");
|
||||
|
||||
try {
|
||||
var connector = tomcat.getConnector();
|
||||
connector.setAsyncTimeout(3000);
|
||||
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 (tomcat != null) {
|
||||
try {
|
||||
tomcat.stop();
|
||||
tomcat.destroy();
|
||||
}
|
||||
catch (LifecycleException e) {
|
||||
throw new RuntimeException("Failed to stop Tomcat", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@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() {
|
||||
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();
|
||||
}
|
||||
|
||||
@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);
|
||||
});
|
||||
|
||||
mcpClient.close();
|
||||
mcpServer.close();
|
||||
}
|
||||
|
||||
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 -> {
|
||||
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 -> {
|
||||
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>> toolsRef = new AtomicReference<>();
|
||||
var mcpClient = clientBuilder.toolsChangeConsumer(toolsUpdate -> {
|
||||
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();
|
||||
toolsRef.set(toolsUpdate);
|
||||
}).sync();
|
||||
|
||||
InitializeResult initResult = mcpClient.initialize();
|
||||
assertThat(initResult).isNotNull();
|
||||
|
||||
assertThat(toolsRef.get()).isNull();
|
||||
|
||||
assertThat(mcpClient.listTools().tools()).contains(tool1.tool());
|
||||
|
||||
mcpServer.notifyToolsListChanged();
|
||||
|
||||
await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
|
||||
assertThat(toolsRef.get()).containsAll(List.of(tool1.tool()));
|
||||
});
|
||||
|
||||
mcpServer.removeTool("tool1");
|
||||
|
||||
await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
|
||||
assertThat(toolsRef.get()).isEmpty();
|
||||
});
|
||||
|
||||
ToolRegistration tool2 = new ToolRegistration(new McpSchema.Tool("tool2", "tool2 description", emptyJsonSchema),
|
||||
request -> callResponse);
|
||||
|
||||
mcpServer.addTool(tool2);
|
||||
|
||||
await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
|
||||
assertThat(toolsRef.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();
|
||||
}
|
||||
|
||||
}
|
||||
Loading…
Reference in a new issue